""" This example is largely adapted from https://github.com/pytorch/examples/blob/master/imagenet/main.py """ from argparse import ArgumentParser, Namespace import os import random from collections import OrderedDict import torch import torch.backends.cudnn as cudnn import torch.nn.functional as F import torch.nn.parallel import torch.optim as optim import torch.optim.lr_scheduler as lr_scheduler import torch.utils.data import torch.utils.data.distributed import torchvision.datasets as datasets import torchvision.models as models import torchvision.transforms as transforms import pytorch_lightning as pl from pytorch_lightning.core import LightningModule # pull out resnet names from torchvision models MODEL_NAMES = sorted( name for name in models.__dict__ if name.islower() and not name.startswith("__") and callable(models.__dict__[name]) ) class ImageNetLightningModel(LightningModule): def __init__(self, arch, pretrained, lr: float, momentum: float, weight_decay: int, data_path: str, batch_size: int, **kwargs): """ TODO: add docstring here """ super().__init__() self.arch = arch self.pretrained = pretrained self.lr = lr self.momentum = momentum self.weight_decay = weight_decay self.data_path = data_path self.batch_size = batch_size self.model = models.__dict__[self.arch](pretrained=self.pretrained) def forward(self, x): return self.model(x) def training_step(self, batch, batch_idx): images, target = batch output = self(images) loss_val = F.cross_entropy(output, target) acc1, acc5 = self.__accuracy(output, target, topk=(1, 5)) tqdm_dict = {'train_loss': loss_val} output = OrderedDict({ 'loss': loss_val, 'acc1': acc1, 'acc5': acc5, 'progress_bar': tqdm_dict, 'log': tqdm_dict }) return output def validation_step(self, batch, batch_idx): images, target = batch output = self(images) loss_val = F.cross_entropy(output, target) acc1, acc5 = self.__accuracy(output, target, topk=(1, 5)) output = OrderedDict({ 'val_loss': loss_val, 'val_acc1': acc1, 'val_acc5': acc5, }) return output def validation_epoch_end(self, outputs): tqdm_dict = {} for metric_name in ["val_loss", "val_acc1", "val_acc5"]: metric_total = 0 for output in outputs: metric_value = output[metric_name] # reduce manually when using dp if self.trainer.use_dp or self.trainer.use_ddp2: metric_value = torch.mean(metric_value) metric_total += metric_value tqdm_dict[metric_name] = metric_total / len(outputs) result = {'progress_bar': tqdm_dict, 'log': tqdm_dict, 'val_loss': tqdm_dict["val_loss"]} return result @classmethod def __accuracy(cls, output, target, topk=(1,)): """Computes the accuracy over the k top predictions for the specified values of k""" with torch.no_grad(): maxk = max(topk) batch_size = target.size(0) _, pred = output.topk(maxk, 1, True, True) pred = pred.t() correct = pred.eq(target.view(1, -1).expand_as(pred)) res = [] for k in topk: correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) res.append(correct_k.mul_(100.0 / batch_size)) return res def configure_optimizers(self): optimizer = optim.SGD( self.parameters(), lr=self.lr, momentum=self.momentum, weight_decay=self.weight_decay ) scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.1) return [optimizer], [scheduler] def train_dataloader(self): normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ) train_dir = os.path.join(self.data_path, 'train') train_dataset = datasets.ImageFolder( train_dir, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if self.use_ddp: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=self.batch_size, shuffle=(train_sampler is None), num_workers=0, sampler=train_sampler ) return train_loader def val_dataloader(self): normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ) val_dir = os.path.join(self.data_path, 'val') val_loader = torch.utils.data.DataLoader( datasets.ImageFolder(val_dir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])), batch_size=self.batch_size, shuffle=False, num_workers=0, ) return val_loader @staticmethod def add_model_specific_args(parent_parser): # pragma: no-cover parser = ArgumentParser(parents=[parent_parser]) parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18', choices=MODEL_NAMES, help='model architecture: ' + ' | '.join(MODEL_NAMES) + ' (default: resnet18)') parser.add_argument('--epochs', default=90, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('--seed', type=int, default=42, help='seed for initializing training. ') parser.add_argument('-b', '--batch-size', default=256, type=int, metavar='N', help='mini-batch size (default: 256), this is the total ' 'batch size of all GPUs on the current node when ' 'using Data Parallel or Distributed Data Parallel') parser.add_argument('--lr', '--learning-rate', default=0.1, type=float, metavar='LR', help='initial learning rate', dest='lr') parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)', dest='weight_decay') parser.add_argument('--pretrained', dest='pretrained', action='store_true', help='use pre-trained model') return parser def get_args(): parent_parser = ArgumentParser(add_help=False) parent_parser.add_argument('--data-path', metavar='DIR', type=str, help='path to dataset') parent_parser.add_argument('--save-path', metavar='DIR', default=".", type=str, help='path to save output') parent_parser.add_argument('--gpus', type=int, default=1, help='how many gpus') parent_parser.add_argument('--distributed-backend', type=str, default='dp', choices=('dp', 'ddp', 'ddp2'), help='supports three options dp, ddp, ddp2') parent_parser.add_argument('--use-16bit', dest='use_16bit', action='store_true', help='if true uses 16 bit precision') parent_parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set') parser = ImageNetLightningModel.add_model_specific_args(parent_parser) return parser.parse_args() def main(args: Namespace) -> None: model = ImageNetLightningModel(**vars(args)) if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True trainer = pl.Trainer( default_root_dir=args.save_path, gpus=args.gpus, max_epochs=args.epochs, distributed_backend=args.distributed_backend, precision=16 if args.use_16bit else 32, ) if args.evaluate: trainer.run_evaluation() else: trainer.fit(model) if __name__ == '__main__': main(get_args())