Spaces:
Sleeping
Sleeping
| # Copyright (c) 2017-present, Facebook, Inc. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the LICENSE file in | |
| # the root directory of this source tree. An additional grant of patent rights | |
| # can be found in the PATENTS file in the same directory. | |
| import os.path as osp | |
| import logging | |
| from dataclasses import dataclass | |
| import torch | |
| from torchvision import transforms | |
| from fairseq.dataclass import FairseqDataclass | |
| from fairseq.tasks import register_task | |
| from fairseq.logging import metrics | |
| try: | |
| from ..data import ImageDataset | |
| except: | |
| import sys | |
| sys.path.append("..") | |
| from data import ImageDataset | |
| from .image_pretraining import ( | |
| ImagePretrainingConfig, | |
| ImagePretrainingTask, | |
| IMG_EXTENSIONS, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| class ImageClassificationConfig(ImagePretrainingConfig): | |
| pass | |
| class ImageClassificationTask(ImagePretrainingTask): | |
| cfg: ImageClassificationConfig | |
| def setup_task(cls, cfg: ImageClassificationConfig, **kwargs): | |
| return cls(cfg) | |
| def load_dataset(self, split: str, task_cfg: FairseqDataclass = None, **kwargs): | |
| data_path = self.cfg.data | |
| cfg = task_cfg or self.cfg | |
| path_with_split = osp.join(data_path, split) | |
| if osp.exists(path_with_split): | |
| data_path = path_with_split | |
| from timm.data import create_transform | |
| if split == "train": | |
| # this should always dispatch to transforms_imagenet_train | |
| transform = create_transform( | |
| input_size=cfg.input_size, | |
| is_training=True, | |
| auto_augment="rand-m9-mstd0.5-inc1", | |
| interpolation="bicubic", | |
| re_prob=0.25, | |
| re_mode="pixel", | |
| re_count=1, | |
| mean=cfg.normalization_mean, | |
| std=cfg.normalization_std, | |
| ) | |
| if not cfg.input_size > 32: | |
| transform.transforms[0] = transforms.RandomCrop( | |
| cfg.input_size, padding=4 | |
| ) | |
| else: | |
| t = [] | |
| if cfg.input_size > 32: | |
| crop_pct = 1 | |
| if cfg.input_size < 384: | |
| crop_pct = 224 / 256 | |
| size = int(cfg.input_size / crop_pct) | |
| t.append( | |
| transforms.Resize( | |
| size, interpolation=3 | |
| ), # to maintain same ratio w.r.t. 224 images | |
| ) | |
| t.append(transforms.CenterCrop(cfg.input_size)) | |
| t.append(transforms.ToTensor()) | |
| t.append( | |
| transforms.Normalize(cfg.normalization_mean, cfg.normalization_std) | |
| ) | |
| transform = transforms.Compose(t) | |
| logger.info(transform) | |
| self.datasets[split] = ImageDataset( | |
| root=data_path, | |
| extensions=IMG_EXTENSIONS, | |
| load_classes=True, | |
| transform=transform, | |
| ) | |
| for k in self.datasets.keys(): | |
| if k != split: | |
| assert self.datasets[k].classes == self.datasets[split].classes | |
| def build_model(self, model_cfg: FairseqDataclass, from_checkpoint=False): | |
| model = super().build_model(model_cfg, from_checkpoint) | |
| actualized_cfg = getattr(model, "cfg", None) | |
| if actualized_cfg is not None: | |
| if hasattr(actualized_cfg, "pretrained_model_args"): | |
| model_cfg.pretrained_model_args = actualized_cfg.pretrained_model_args | |
| return model | |
| def reduce_metrics(self, logging_outputs, criterion): | |
| super().reduce_metrics(logging_outputs, criterion) | |
| if "correct" in logging_outputs[0]: | |
| zero = torch.scalar_tensor(0.0) | |
| correct = sum(log.get("correct", zero) for log in logging_outputs) | |
| metrics.log_scalar_sum("_correct", correct) | |
| metrics.log_derived( | |
| "accuracy", | |
| lambda meters: 100 * meters["_correct"].sum / meters["sample_size"].sum, | |
| ) | |