def __getitem__(self, index): if self.is_train: img, target = imageio.imread( self.train_img_label[index][0]), self.train_img_label[index][1] if len(img.shape) == 2: img = np.stack([img] * 3, 2) img = Image.fromarray(img, mode='RGB') img = transforms.Resize((self.input_size, self.input_size), Image.BILINEAR)(img) if (rand_aug): img = RandAugment(N, M)(img) # img = transforms.RandomResizedCrop(size=self.input_size,scale=(0.4, 0.75),ratio=(0.5,1.5))(img) # img = transforms.RandomCrop(self.input_size)(img) img = transforms.RandomHorizontalFlip()(img) img = transforms.ColorJitter(brightness=0.2, contrast=0.2)(img) img = transforms.ToTensor()(img) img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img) else: img, target = imageio.imread( self.test_img_label[index][0]), self.test_img_label[index][1] if len(img.shape) == 2: img = np.stack([img] * 3, 2) img = Image.fromarray(img, mode='RGB') img = transforms.Resize((self.input_size, self.input_size), Image.BILINEAR)(img) # img = transforms.CenterCrop(self.input_size)(img) img = transforms.ToTensor()(img) img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img) return img, target
def randagument(self): self.transform = transforms.Compose([ transforms.CenterCrop(self.size), RandAugment(1, 15), ToTensor(), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ])
def UnlabeledImageDataLoader(transform_args, randaug = False, image_dir = '../../ChestX-ray14/images/', label_dir_list = ['../../ChestX-ray14/train_val_list.txt', '../../ChestX-ray14/test_list.txt'], batchsize = 64, numworker = 12, res = 456, train = True, ): # --- Transforms --- transform_list = [getattr(transforms, transform_args[0])(res), *[getattr(transforms, single_transform)() for single_transform in transform_args[1:]]] img_transform = transforms.Compose(transform_list) if train and randaug: img_transform.transforms.insert(0, RandAugment(2, 3)) # --- Data Collection --- unlabeled_dataset = UnlabeledImageDataset(image_dir=image_dir, label_dir_list=label_dir_list, transform=img_transform) # --- DataLoader --- unlabeled_loader = torch.utils.data.DataLoader(unlabeled_dataset, batch_size=batchsize, num_workers=numworker, pin_memory=True) return unlabeled_loader, unlabeled_dataset.label_dir
def randaugment(self): self.transform = transforms.Compose([ transforms.CenterCrop(self.size), transforms.RandomHorizontalFlip(p=0.5), RandAugment(1, 15), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ])
def get_simclr_data_transforms_randAugment(input_shape): # get a set of data augmentation transformations as described in the SimCLR paper. data_transforms = transforms.Compose([ transforms.RandomResizedCrop(size=eval(input_shape)[0]), transforms.RandomHorizontalFlip(), transforms.ToTensor() ]) data_transforms.transforms.insert(0, RandAugment(3, 9)) # pip install git+https://github.com/ildoonet/pytorch-randaugment return data_transforms
def __init__(self, name, data, targets, transforms, base_transforms): self.name = name self.data = data self.targets = targets self.transforms = transforms self.base_transforms = base_transforms self.rand_aug_transforms = copy.deepcopy(self.base_transforms) self.committee_size = 1 self.ra_obj = RandAugment(1, 2.0) if self.name == 'mnist': # See if this can be dropped entirely without losing performance self.transforms.transforms.insert(0, torchvision.transforms.Grayscale(num_output_channels=1)) self.base_transforms.transforms.insert(0, torchvision.transforms.Grayscale(num_output_channels=1)) self.rand_aug_transforms.transforms.insert(0, torchvision.transforms.Grayscale(num_output_channels=1)) self.rand_aug_transforms.transforms.insert(0, self.ra_obj)
def prepare_data(self): ##### RandAugment changes here for now data_transforms = { # same for now 'train': transforms.Compose([ transforms.Resize((256, 256), interpolation=2), # transforms.RandomCrop(32, padding=0), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[n / 255. for n in [129.3, 124.1, 112.4]], std=[n / 255. for n in [68.2, 65.4, 70.4]]) ]), 'valid': transforms.Compose([ transforms.Resize((256, 256), interpolation=2), # transforms.RandomCrop(32, padding=0), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[n / 255. for n in [129.3, 124.1, 112.4]], std=[n / 255. for n in [68.2, 65.4, 70.4]]) ]), } data_transforms['train'].transforms.insert( 0, RandAugment(self.randAugParamN, self.randAugParamM)) train_path = '/lab/vislab/OPEN/datasets_RGB_new/train/' valid_path = '/lab/vislab/OPEN/datasets_RGB_new/val/' self.trainset = torchvision.datasets.ImageFolder( train_path, data_transforms['train']) self.validset = torchvision.datasets.ImageFolder( valid_path, data_transforms['valid'])
def ImageDataLoader(transform_args, pseudo_label_args, randaug = False, image_dir = '../../', label_dir = '../../CheXpert-v1.0-small/train.csv', batchsize = 64, numworker = 12, res = 456, train = True, label_smooth = False, ): """ train_dataset (images): [0]: Minibatch of Images. torch.Tensor of size (batchsize, 3, H, W) [1]: What subfolder the data come from. list of size (batchsize) train_label: 1: Positive 0: Negative + Unmentioned -1: Uncertain """ # --- Transforms --- transform_list = [getattr(transforms, transform_args[0])(res), *[getattr(transforms, single_transform)() for single_transform in transform_args[1:]]] img_transform = transforms.Compose(transform_list) if train and randaug: img_transform.transforms.insert(0, RandAugment(2, 3)) # --- Data Collection --- train_dataset = ImageDataset(image_dir=image_dir, label_dir=label_dir, transform=img_transform, train=train, label_smooth=label_smooth) if pseudo_label_args['use_pseudo_label']: ssl_dataset = ImageDataset(image_dir=pseudo_label_args['pseudo_dataset_args']['image_dir'], label_dir=pseudo_label_args['pseudo_dataset_args']['label_dir'], transform=img_transform, train=train, label_smooth=label_smooth) train_dataset = torch.utils.data.ConcatDataset((train_dataset, ssl_dataset)) # --- DataLoader --- train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batchsize, num_workers=numworker, pin_memory=True) return train_loader, len(train_dataset)
def load_dataset(N, M, dataset, dataset_dir): '''Add RandAugment with N, M(hyperparameter) N: Number of augmentation transformations to apply sequentially. M: Magnitude for all the transformations.''' transform_randaugment = transforms.Compose([ RandAugment(N, M), #transforms.Resize(32), transforms.ToTensor(), ]) transform_resize = transforms.Compose([ #transforms.Resize(32), transforms.ToTensor(), ]) if dataset == 'stl10': stl10 = {} stl10['train_augmented'] = STL10(dataset_dir, 'train', transform=transform_randaugment, target_transform=one_hot, download=True) stl10['train_unaugmented'] = STL10(dataset_dir, 'train', transform=transform_resize, target_transform=one_hot) stl10['unlabeled_augmented'] = STL10(dataset_dir, 'unlabeled', transform=transform_randaugment, download=True) stl10['unlabeled_unaugmented'] = STL10(dataset_dir, 'unlabeled', transform=transform_resize, download=True) stl10['test'] = STL10(dataset_dir, 'test', transform=transform_resize, download=True) return stl10 else: raise Excpetion(f"Dataset '{dataset}' not implemented")
def get_dataset(dataset, cutout_length=0, N=3, M=5, RandA=False): MEAN = [0.49139968, 0.48215827, 0.44653124] STD = [0.24703233, 0.24348505, 0.26158768] transf = [ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip() ] normalize = [ transforms.Resize(32), transforms.ToTensor(), transforms.Normalize(MEAN, STD) ] cutout = [] if cutout_length > 0: cutout.append(Cutout(cutout_length)) train_transform = transforms.Compose(transf + normalize + cutout) if RandA: # Add RandAugment with N, M(hyperparameter) train_transform.transforms.insert(0, RandAugment(N, M)) valid_transform = transforms.Compose(normalize) if dataset == "cifar10": dataset_train = CIFAR10(root="./data", train=True, download=True, transform=train_transform) dataset_valid = CIFAR10(root="./data", train=False, download=True, transform=valid_transform) elif dataset == 'cifar100': dataset_train = CIFAR100(root="./data", train=True, download=True, transform=train_transform) dataset_valid = CIFAR100(root="./data", train=False, download=True, transform=valid_transform) else: raise NotImplementedError return dataset_train, dataset_valid
import argparse def parse_args(): parser = argparse.ArgumentParser( description = "Arguments for ModelNet Visualization", formatter_class = argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument("-N", type=int, default=4, help="RandAugment N") parser.add_argument("-M", type=int, default=4, help="RandAugment M") return parser.parse_args() args = parse_args() transforms = transforms.Compose( [ RandAugment(n=args.N, m=args.M), d_utils.PointCloudToTensor(), d_utils.PointCloudNormalize(), ] ) dset = ModelNet40Cls(1024, train = True, transforms=transforms) pc , _ = dset[100] pc = pc[:,0:3] #dloader = torch.utils.data.DataLoader(dset, batch_size = 1, shuffle=True) xyz = pc.numpy() # Pass xyz to Open3D.o3d.geometry.PointCloud and visualize pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(xyz) o3d.io.write_point_cloud("sync.ply", pcd)
def main_worker(gpu, ngpus_per_node, args): CHECKPOINT_ID = "{}_{}epochs_{}bsz_{:0.4f}lr" \ .format(args.id[:5], args.epochs, args.batch_size, args.lr) if args.mlp: CHECKPOINT_ID += "_mlp" if args.aug_plus: CHECKPOINT_ID += "_augplus" if args.cos: CHECKPOINT_ID += "_cos" if args.faa_aug: CHECKPOINT_ID += "_faa" if args.randomcrop: CHECKPOINT_ID += "_randcrop" if args.rotnet: CHECKPOINT_ID += "_rotnet" if args.rand_aug: CHECKPOINT_ID += "_randaug" if not (args.kfold == None): CHECKPOINT_ID += "_fold_%d" % (args.kfold) if not (args.custom_aug_name == None): CHECKPOINT_ID += "_custom_aug_" + args.custom_aug_name CHECKPOINT_ID += args.dataid args.gpu = gpu # suppress printing if not master if args.multiprocessing_distributed and args.gpu != 0: def print_pass(*args): pass builtins.print = print_pass if args.gpu is not None: print("Use GPU: {} for training".format(args.gpu)) if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model heads = {} if not args.nomoco: heads["moco"] = {"num_classes": args.moco_dim} if args.rotnet: heads["rotnet"] = {"num_classes": 4} model = moco.builder.MoCo(models.__dict__[args.arch], K=args.moco_k, m=args.moco_m, T=args.moco_t, mlp=args.mlp, dataid=args.dataid, multitask_heads=heads) print(model) # setup file structure for saving pathlib.Path(args.checkpoint_fp).mkdir(parents=True, exist_ok=True) if args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int( (args.workers + ngpus_per_node - 1) / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.gpu]) else: model.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set model = torch.nn.parallel.DistributedDataParallel(model) elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) # comment out the following line for debugging # raise NotImplementedError("Only DistributedDataParallel is supported.") else: # AllGather implementation (batch shuffle, queue update, etc.) in # this code only supports DistributedDataParallel. raise NotImplementedError("Only DistributedDataParallel is supported.") # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda(args.gpu) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) if args.gpu is None: checkpoint = torch.load(args.resume) else: # Map model to be loaded to specified single gpu. loc = 'cuda:{}'.format(args.gpu) checkpoint = torch.load(args.resume, map_location=loc) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) args.id = checkpoint['id'] args.name = checkpoint['name'] CHECKPOINT_ID = checkpoint['name'] print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code # Set up crops and normalization depending on the dataset. # Cifar 10 crops and normalization. if args.dataid == "cifar10" or args.dataid == "svhn": _CIFAR_MEAN, _CIFAR_STD = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010) normalize = transforms.Normalize(mean=_CIFAR_MEAN, std=_CIFAR_STD) if not args.randomcrop: random_resized_crop = transforms.RandomResizedCrop( 28, scale=(args.rrc_param, 1.)) else: # Use the crop they were using in Fast AutoAugment. random_resized_crop = transforms.RandomCrop(32, padding=4) # Use the imagenet parameters. elif args.dataid == "imagenet" or args.dataid == "logos": normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) random_resized_crop = transforms.RandomResizedCrop(224, scale=(0.2, 1.)) if args.aug_plus and (args.faa_aug or args.rand_aug or args.rand_aug_orig or not (args.custom_aug_name == None)): raise Exception("Cannot have multiple augs on command line") if args.aug_plus: # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709 augmentation = [ random_resized_crop, transforms.RandomApply( [ transforms.ColorJitter(0.4, 0.4, 0.4, 0.1) # not strengthened ], p=0.8), transforms.RandomGrayscale(p=0.2), transforms.RandomApply( [moco.loader.GaussianBlur([args.sigma / 20, args.sigma])], p=0.5), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ] elif args.faa_aug: augmentation, _ = slm_utils.get_faa_transforms.get_faa_transforms_cifar_10( args.randomcrop, args.gauss) transformations = moco.loader.TwoCropsTransform(augmentation) elif args.rand_aug_orig: print("Using random aug original") augmentation = [ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), RandAugment(args.rand_aug_n, args.rand_aug_m), transforms.ToTensor(), normalize ] elif args.rand_aug: randaug_n = args.rand_aug_n if args.rand_aug_linear_m: print("Using random aug with linear m") randaug_m = args.rand_aug_m_min else: randaug_m = args.rand_aug_m print("Using random aug") if args.rand_aug_top_k > 0: randaug = TopRandAugment(randaug_n, randaug_m, args.rand_aug_top_k) else: randaug = RandAugment(randaug_n, randaug_m) augmentation = [ random_resized_crop, transforms.RandomHorizontalFlip(), randaug, transforms.ToTensor(), normalize ] elif args.rand_resize_only and args.custom_aug_name == None: print("Using random resize only") augmentation = [ random_resized_crop, transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ] elif not args.custom_aug_name == None: augmentation, _ = slm_utils.get_faa_transforms.load_custom_transforms( name=args.custom_aug_name, randomcrop=args.randomcrop, aug_idx=args.single_aug_idx, dataid=args.dataid) print('using custom augs', augmentation) transformations = moco.loader.TwoCropsTransform(augmentation) else: # MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978 print('using v1 augs') augmentation = [ random_resized_crop, transforms.RandomGrayscale(p=0.2), transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ] if not args.faa_aug and args.custom_aug_name == None: print('using augmentation', augmentation) transformations = moco.loader.TwoCropsTransform( transforms.Compose(augmentation)) print('xforms', transformations) if args.dataid == "imagenet" and not args.reduced_imgnet: train_dataset = datasets.ImageFolder(args.data, transformations) elif args.dataid == "logos" and not args.reduced_imgnet: train_dataset = data_loader.GetLoader( data_root=args.data, data_list='train_images_root.txt', transform=transformations) elif (args.dataid == "imagenet" or args.dataid == 'logos') and args.reduced_imgnet: # idx120 = [16, 23, 52, 57, 76, 93, 95, 96, 99, 121, 122, 128, 148, 172, 181, 189, 202, 210, 232, 238, 257, 258, 259, 277, 283, 289, 295, 304, 307, 318, 322, 331, 337, 338, 345, 350, 361, 375, 376, 381, 388, 399, 401, 408, 424, 431, 432, 440, 447, 462, 464, 472, 483, 497, 506, 512, 530, 541, 553, 554, 557, 564, 570, 584, 612, 614, 619, 626, 631, 632, 650, 657, 658, 660, 674, 675, 680, 682, 691, 695, 699, 711, 734, 736, 741, 754, 757, 764, 769, 770, 780, 781, 787, 797, 799, 811, 822, 829, 830, 835, 837, 842, 843, 845, 873, 883, 897, 900, 902, 905, 913, 920, 925, 937, 938, 940, 941, 944, 949, 959] if args.dataid == "imagenet": total_trainset = ImageNet( root=args.data, transform=transformations ) # TODO for LINCLS, make this train and test xforms. else: total_trainset = data_loader.GetLoader( data_root=args.data, data_list='train_images_root.txt', transform=transformations) train_idx = np.arange(len(total_trainset)) np.random.seed(1337) #fingers crossed. np.random.shuffle(train_idx) train_idx = train_idx[:50000] kfold = args.kfold print('KFOLD BEING USED', kfold) subset = np.arange(kfold * 10000, (kfold + 1) * 10000) print('start', 'end', kfold * 10000, (kfold + 1) * 10000) valid_idx = train_idx[subset] train_idx = np.delete(train_idx, subset) print('first val_idx', valid_idx[:10]) train_dataset = total_trainset train_dataset = Subset(train_dataset, train_idx) train_sampler = SubsetRandomSampler(train_idx) valid_sampler = SubsetSampler(valid_idx) print(len(train_dataset)) train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) print(len(train_dataset)) print('first 10 train', train_idx[:10]) print('first 10 valid', valid_idx[:10]) print('len train', len(train_idx)) print('len valid', len(valid_idx)) for i in valid_idx: if i in train_idx: raise Exception("Valid idx in train idx: this is unexpected") print('train_sampler', train_sampler) elif args.dataid == "cifar10": train_dataset = torchvision.datasets.CIFAR10(args.data, transform=transformations, download=True) elif args.dataid == "svhn": train_dataset = torchvision.datasets.SVHN(args.data, transform=transformations, download=True) else: raise NotImplementedError( "Support for the following dataset is not yet implemented: {}". format(args.dataid)) if not args.kfold == None and not args.reduced_imgnet: torch.manual_seed(1337) print('before: K FOLD', args.kfold, len(train_dataset)) lengths = [len(train_dataset) // 5] * 5 print(lengths) lengths[-1] = int(lengths[-1] + (len(train_dataset) - np.sum(lengths))) print(lengths) folds = torch.utils.data.random_split(train_dataset, lengths) print(len(folds)) folds.pop(args.kfold) print(len(folds)) train_dataset = torch.utils.data.ConcatDataset(folds) print(len(train_dataset)) else: print("NO KFOLD ARG", args.kfold, ' or ', args.reduced_imgnet) if args.distributed and not args.reduced_imgnet: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) elif not args.reduced_imgnet: train_sampler = None print('train sampler', train_sampler) torch.manual_seed(1337) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True) print(len(train_loader)) # CR: only the master will report to wandb for now if not args.multiprocessing_distributed or args.rank % ngpus_per_node == 0: wandb.init(project=args.wandbproj, name=CHECKPOINT_ID, id=args.id, resume=args.resume, config=args.__dict__, notes=args.notes) print(model) for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch, args) if args.rand_aug_linear_m: mval = args.rand_aug_m_min + math.floor( float(epoch) / float(args.epochs) * (args.rand_aug_m_max - args.rand_aug_m_min + 1)) print("Rand aug m: {}".format(mval)) randaug.m = mval # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args, CHECKPOINT_ID) # save current epoch if not args.nosave_latest and (not args.multiprocessing_distributed or args.rank % ngpus_per_node == 0): print("saving latest epoch") cp_filename = "{}_latest.tar".format(CHECKPOINT_ID[:5]) cp_fullpath = os.path.join(args.checkpoint_fp, cp_filename) torch.save( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'id': args.id, 'name': CHECKPOINT_ID, }, cp_fullpath) print("saved latest epoch") if (epoch % args.checkpoint_interval == 0 or epoch == args.epochs-1) \ and (not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0)): cp_filename = "{}_{:04d}.tar".format(CHECKPOINT_ID, epoch) cp_fullpath = os.path.join(args.checkpoint_fp, cp_filename) torch.save( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'id': args.id, 'name': CHECKPOINT_ID, }, cp_fullpath) if args.upload_checkpoints: print("Uploading wandb checkpoint") wandb.save(cp_fullpath) if epoch == args.epochs - 1: print("Saving final results to wandb") wandb.save(cp_fullpath) print("Done - wrapping up")
def configure_transform(phase: str, transform_type: str): if phase == "train": if transform_type == Aug.BaseTransform: transform = transforms.Compose([ transforms.Resize((512, 384), Image.BILINEAR), transforms.CenterCrop((384, 384)), transforms.RandomResizedCrop((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) elif transform_type == Aug.FaceCrop: transform = transforms.Compose([ transforms.Lambda(lambda x: crop(x)), transforms.Resize((312, 234)), transforms.RandomCrop((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) elif transform_type == Aug.Random: transform = transforms.Compose([ transforms.CenterCrop((384, 384)), transforms.RandomResizedCrop((224, 224)), RandAugment(2, 9), # N: 몇 개 선택할지 M: 몇 번 변화시킬 것인지 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) elif transform_type == Aug.TTA: # 'random'과 동일 transform = transforms.Compose([ transforms.CenterCrop((384, 384)), transforms.RandomResizedCrop((224, 224)), RandAugment(2, 9), # N: 몇 개 선택할지 M: 몇 번 변화시킬 것인지 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) else: raise NotImplementedError() else: if transform_type == Aug.BaseTransform: transform = transforms.Compose([ transforms.Resize((512, 384), Image.BILINEAR), transforms.CenterCrop((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) elif transform_type == Aug.FaceCrop: transform = transforms.Compose([ transforms.Lambda(lambda x: crop(x)), transforms.Resize((312, 234)), transforms.CenterCrop((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) elif transform_type == Aug.Random: transform = transforms.Compose([ transforms.CenterCrop((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) elif transform_type == Aug.TTA: transform = transforms.Compose([ transforms.CenterCrop((384, 384)), transforms.RandomResizedCrop((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) else: raise NotImplementedError() return transform
# Augmentations. transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(cifar10_mean, cifar10_std) ]) transform_strong = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(cifar10_mean, cifar10_std) ]) transform_strong.transforms.insert(0, RandAugment(3, 4)) transform_strong.transforms.append(CutoutDefault(16)) transform_val = transforms.Compose( [transforms.ToTensor(), transforms.Normalize(cifar10_mean, cifar10_std)]) class TransformMix: def __init__(self, transform): self.transform = transform def __call__(self, inp): out1 = self.transform(inp) out2 = self.transform(inp) return out1, out2
def main(): args = parse_args() if args.name is None: args.name = '%s_%s' % (args.arch, datetime.now().strftime('%m%d%H%M')) if not os.path.exists('models/%s' % args.name): os.makedirs('models/%s' % args.name) if args.print_model: print('Config -----') for arg in vars(args): print('- %s: %s' % (arg, getattr(args, arg))) print('------------') with open('models/%s/args.txt' % args.name, 'w') as f: for arg in vars(args): print('- %s: %s' % (arg, getattr(args, arg)), file=f) joblib.dump(args, 'models/%s/args.pkl' % args.name) if args.loss == 'CrossEntropyLoss': criterion = nn.CrossEntropyLoss().cuda() elif args.loss == 'FocalLoss': criterion = FocalLoss().cuda() elif args.loss == 'MSELoss': criterion = nn.MSELoss().cuda() elif args.loss == 'BCELoss': criterion = nn.BCELoss().cuda() elif args.loss == 'L1Loss': criterion = nn.L1Loss().cuda() elif args.loss == 'multitask': criterion = { 'classification': nn.CrossEntropyLoss().cuda(), 'regression': nn.MSELoss().cuda(), } else: raise NotImplementedError if args.pred_type == 'classification': num_outputs = 5 elif args.pred_type == 'regression': num_outputs = 1 elif args.loss == 'multitask': num_outputs = 6 else: raise NotImplementedError cudnn.benchmark = True train_transform = [] if args.RAugment: train_transform = transforms.Compose([ transforms.Resize((args.img_size, args.img_size)), transforms.RandomAffine( degrees=(args.rotate_min, args.rotate_max) if args.rotate else 0, translate=(args.translate_min, args.translate_max) if args.translate else None, scale=(args.rescale_min, args.rescale_max) if args.rescale else None, shear=(args.shear_min, args.shear_max) if args.shear else None, ), transforms.CenterCrop(args.input_size), transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0), # transforms.RandomVerticalFlip(p=0.5 if args.flip else 0), transforms.ColorJitter( brightness=(args.brightness_min, args.brightness_max) if args.brightness else 0, contrast=(args.contrast_min, args.contrast_max) if args.contrast else 0, saturation=(args.saturation_min, args.saturation_max) if args.saturation else 0, hue=0), RandomErase( prob=args.random_erase_prob if args.random_erase else 0, sl=args.random_erase_sl, sh=args.random_erase_sh, r=args.random_erase_r), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) train_transform.transforms.insert( 0, RandAugment(args.RAugment_N, args.RAugment_M)) else: train_transform = transforms.Compose([ transforms.Resize((args.img_size, args.img_size)), transforms.RandomAffine( degrees=(args.rotate_min, args.rotate_max) if args.rotate else 0, translate=(args.translate_min, args.translate_max) if args.translate else None, scale=(args.rescale_min, args.rescale_max) if args.rescale else None, shear=(args.shear_min, args.shear_max) if args.shear else None, ), transforms.CenterCrop(args.input_size), transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0), # transforms.RandomVerticalFlip(p=0.5 if args.flip else 0), transforms.ColorJitter( brightness=(args.brightness_min, args.brightness_max) if args.brightness else 0, contrast=(args.contrast_min, args.contrast_max) if args.contrast else 0, saturation=(args.saturation_min, args.saturation_max) if args.saturation else 0, hue=0), RandomErase( prob=args.random_erase_prob if args.random_erase else 0, sl=args.random_erase_sl, sh=args.random_erase_sh, r=args.random_erase_r), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) # IF Class Aware, need to do Data Augumentation on Validation. if args.class_aware: val_transform = transforms.Compose([ transforms.Resize((args.img_size, args.input_size)), transforms.RandomAffine( degrees=(args.rotate_min, args.rotate_max) if args.rotate else 0, translate=(args.translate_min, args.translate_max) if args.translate else None, scale=(args.rescale_min, args.rescale_max) if args.rescale else None, shear=(args.shear_min, args.shear_max) if args.shear else None, ), transforms.CenterCrop(args.input_size), transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0), # transforms.RandomVerticalFlip(p=0.5 if args.flip else 0), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) else: val_transform = transforms.Compose([ transforms.Resize((args.img_size, args.input_size)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) # Data loading. if 'diabetic_retinopathy' in args.train_dataset: diabetic_retinopathy_dir = preprocess('diabetic_retinopathy', args.img_size, scale=args.scale_radius, norm=args.normalize, pad=args.padding, remove=args.remove) diabetic_retinopathy_df = pd.read_csv( 'inputs/diabetic-retinopathy-resized/trainLabels.csv') diabetic_retinopathy_img_paths = \ diabetic_retinopathy_dir + '/' + diabetic_retinopathy_df['image'].values + '.jpeg' diabetic_retinopathy_labels = diabetic_retinopathy_df['level'].values if 'RNFL' in args.train_dataset: aptos2019_dir = preprocess('RNFL', args.img_size, scale=args.scale_radius, norm=args.normalize, pad=args.padding, remove=args.remove, red_free=args.red_free) aptos2019_df = pd.read_csv(Path("./data/dataset/train/label.csv")) aptos2019_img_paths = aptos2019_dir + '/' + aptos2019_df[ 'id_code'].values + '.png' aptos2019_labels = aptos2019_df['diagnosis'].values if args.train_dataset == 'RNFL': skf = StratifiedKFold(n_splits=args.n_splits, shuffle=True, random_state=41) img_paths = [] labels = [] for fold, (train_idx, val_idx) in enumerate( skf.split(aptos2019_img_paths, aptos2019_labels)): img_paths.append( (aptos2019_img_paths[train_idx], aptos2019_img_paths[val_idx])) labels.append( (aptos2019_labels[train_idx], aptos2019_labels[val_idx])) elif args.train_dataset == 'diabetic_retinopathy': img_paths = [(diabetic_retinopathy_img_paths, aptos2019_img_paths)] labels = [(diabetic_retinopathy_labels, aptos2019_labels)] elif 'diabetic_retinopathy' in args.train_dataset and 'aptos2019' in args.train_dataset: skf = StratifiedKFold(n_splits=args.n_splits, shuffle=True, random_state=41) img_paths = [] labels = [] for fold, (train_idx, val_idx) in enumerate( skf.split(aptos2019_img_paths, aptos2019_labels)): img_paths.append((np.hstack((aptos2019_img_paths[train_idx], diabetic_retinopathy_img_paths)), aptos2019_img_paths[val_idx])) labels.append((np.hstack( (aptos2019_labels[train_idx], diabetic_retinopathy_labels)), aptos2019_labels[val_idx])) # else: # raise NotImplementedError if args.pseudo_labels: test_df = pd.read_csv('probs/%s.csv' % args.pseudo_labels) test_dir = preprocess('test', args.img_size, scale=args.scale_radius, norm=args.normalize, pad=args.padding, remove=args.remove) test_img_paths = test_dir + '/' + test_df['id_code'].values + '.png' test_labels = test_df['diagnosis'].values for fold in range(len(img_paths)): img_paths[fold] = (np.hstack( (img_paths[fold][0], test_img_paths)), img_paths[fold][1]) labels[fold] = (np.hstack( (labels[fold][0], test_labels)), labels[fold][1]) if 'messidor' in args.train_dataset: test_dir = preprocess('messidor', args.img_size, scale=args.scale_radius, norm=args.normalize, pad=args.padding, remove=args.remove) folds = [] best_losses = [] best_scores = [] for fold, ((train_img_paths, val_img_paths), (train_labels, val_labels)) in enumerate(zip(img_paths, labels)): print('Fold [%d/%d]' % (fold + 1, len(img_paths))) if os.path.exists('models/%s/model_%d.pth' % (args.name, fold + 1)): log = pd.read_csv('models/%s/log_%d.csv' % (args.name, fold + 1)) best_loss, best_score = log.loc[log['val_loss'].values.argmin(), ['val_loss', 'val_score']].values folds.append(str(fold + 1)) best_losses.append(best_loss) best_scores.append(best_score) continue if args.remove_duplicate: md5_df = pd.read_csv('inputs/strMd5.csv') duplicate_img_paths = aptos2019_dir + '/' + md5_df[ (md5_df.strMd5_count > 1) & (~md5_df.diagnosis.isnull())]['id_code'].values + '.png' print(duplicate_img_paths) for duplicate_img_path in duplicate_img_paths: train_labels = train_labels[ train_img_paths != duplicate_img_path] train_img_paths = train_img_paths[ train_img_paths != duplicate_img_path] val_labels = val_labels[val_img_paths != duplicate_img_path] val_img_paths = val_img_paths[ val_img_paths != duplicate_img_path] # Train. train_set = Dataset(train_img_paths, train_labels, transform=train_transform) _, class_sample_counts = np.unique(train_labels, return_counts=True) _, vclass_sample_counts = np.unique(val_labels, return_counts=True) print('Train:', class_sample_counts) print('Validation:', vclass_sample_counts) if args.class_aware: print('Class Aware(N:P):', args.N_ratio, ':', args.P_ratio) print('Number of Train samples:', args.num_train_sample) print('Number of Validation samples:', args.num_val_sample) print('------------------------------------------') weights = 1. / torch.tensor(class_sample_counts, dtype=torch.float) weights = np.array([1, 1, 1, 1, 1]) samples_weights = weights[train_labels] vsamples_weights = weights[val_labels] tsampler = WeightedRandomSampler(weights=samples_weights, num_samples=args.num_train_sample, replacement=True) vsampler = WeightedRandomSampler(weights=vsamples_weights, num_samples=args.num_val_sample, replacement=True) train_loader = torch.utils.data.DataLoader( train_set, batch_size=args.batch_size, shuffle=False if args.class_aware else True, num_workers=args.num_workers, sampler=tsampler if args.class_aware else None) val_set = Dataset(val_img_paths, val_labels, transform=val_transform) val_loader = torch.utils.data.DataLoader( val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, sampler=vsampler if args.class_aware else None) # Create model. model = get_model(model_name=args.arch, num_outputs=num_outputs, freeze_bn=args.freeze_bn, dropout_p=args.dropout_p) if args.print_model: print(model) model = model.cuda() if args.pretrained_model is not None: model.load_state_dict( torch.load('models/%s/model_%d.pth' % (args.pretrained_model, fold + 1))) # Print the model. if args.optimizer == 'Adam': optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'AdamW': optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'RAdam': optimizer = RAdam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'SGD': optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) elif args.optimizer == 'Ranger': optimizer = Ranger(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'RangerLars': optimizer = RangerLars(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'SAM': base_optimizer = RangerLars optimizer = SAM(model.parameters(), base_optimizer, lr=args.lr, weight_decay=args.weight_decay) # if args.scheduler == 'CosineAnnealingLR': # scheduler = lr_scheduler.CosineAnnealingLR( # optimizer, T_max=args.epochs, eta_min=args.min_lr) if args.scheduler == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=5, T_mult=6, eta_min=args.min_lr) elif args.scheduler == 'ReduceLROnPlateau': scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=args.factor, patience=args.patience, verbose=1, min_lr=args.min_lr) log = pd.DataFrame( index=[], columns=['epoch', 'loss', 'score', 'val_loss', 'val_score']) log = { 'epoch': [], 'loss': [], 'score': [], 'val_loss': [], 'val_score': [], } best_loss = float('inf') best_score = 0 for epoch in range(args.epochs): print('Epoch [%d/%d]' % (epoch + 1, args.epochs)) # Train for one epoch. train_loss, train_score = train(args, train_loader, model, criterion, optimizer, epoch) # Evaluate on validation set. val_loss, val_score = validate(args, val_loader, model, criterion) if args.scheduler == 'CosineAnnealingLR': scheduler.step() elif args.scheduler == 'ReduceLROnPlateau': scheduler.step(val_loss) print( 'loss %.4f - score %.4f - val_loss %.4f - val_score %.4f - best_score %.4f' % (train_loss, train_score, val_loss, val_score, best_score)) log['epoch'].append(epoch) log['loss'].append(train_loss) log['score'].append(train_score) log['val_loss'].append(val_loss) log['val_score'].append(val_score) pd.DataFrame(log).to_csv('models/%s/log_%d.csv' % (args.name, fold + 1), index=False) if val_score >= best_score: if (val_score / train_score) <= 1.6: torch.save( model.state_dict(), 'models/%s/model_%d.pth' % (args.name, fold + 1)) best_loss = val_loss best_score = val_score print("=> saved best model") print('val_loss: %f' % best_loss) print('val_score: %f' % best_score) folds.append(str(fold + 1)) best_losses.append(best_loss) best_scores.append(best_score) results = pd.DataFrame({ 'fold': folds + ['mean'], 'best_loss': best_losses + [np.mean(best_losses)], 'best_score': best_scores + [np.mean(best_scores)], }) print(results) results.to_csv('models/%s/results.csv' % args.name, index=False) torch.cuda.empty_cache() if not args.cv: break
def get(conf, trial, param_gen): SIZE = conf["input_size"] # TRANSFORM trans = [ transforms.Resize((SIZE, SIZE)), transforms.CenterCrop(SIZE), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ] trans = transforms.Compose(trans) if conf["use_augmentation"]: train_trans = [ transforms.Resize((SIZE, SIZE)), transforms.CenterCrop(SIZE) ] # Flip if param_gen.suggest_categorical("augment_flip", [0, 1]) > 0: train_trans.append(transforms.RandomHorizontalFlip()) # Rotation, Shift, Scale rot_degree = param_gen.suggest_uniform("augment_rot", 0.0, 180.0) translate = param_gen.suggest_uniform("augment_trans", 0.0, 0.3) scale = param_gen.suggest_uniform("augment_scale", 1.0, 1.3) train_trans.append( transforms.RandomAffine(rot_degree, translate=(translate, translate), scale=(1 / scale, scale))) train_trans += [ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ] # Erasing if param_gen.suggest_categorical("augment_erase", [0, 1]) > 0: train_trans.append(transforms.RandomErasing()) train_trans = transforms.Compose(train_trans) elif conf["use_randaug"]: _N = param_gen.suggest_int("randaug_N", 0, 10) _M = param_gen.suggest_int("randaug_M", 0, 30) gray_p = param_gen.suggest_uniform("gray_p", 0.0, 1.0) train_trans = [ transforms.Resize((SIZE, SIZE)), transforms.CenterCrop(SIZE), RandAugment(_N, _M), transforms.RandomGrayscale(gray_p), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ] train_trans = transforms.Compose(train_trans) else: train_trans = trans trans_dict = {"train": train_trans, "valid": trans} # DATASET GENERATOR label_dirnames = glob(os.path.join(conf["data_dir"], "*")) labels = [] for dirname in label_dirnames: label = dirname.split(os.path.sep)[-1] labels.append(label) labels = sorted(labels) def dataset_fold_generator(all=False): folds = KFold(n_splits=conf["n_fold"], shuffle=True, random_state=1234) dummy = np.arange(len(labels), dtype=np.int32)[:, np.newaxis] for train_lab, dev_lab in list( folds.split(dummy))[:conf["fold_trials"]]: train_lab = set([labels[_] for _ in train_lab]) dev_lab = set([labels[_] for _ in dev_lab]) train_dataset = ImageFolder(root=conf["data_dir"], transform=trans_dict["train"], is_valid_file=lambda fn: fn.split( os.path.sep)[-2] in train_lab) train_sampler = samplers.MPerClassSampler( labels=train_dataset.targets, m=param_gen.suggest_categorical("m", conf["m_cands"]), batch_size=conf["batch_size"], length_before_new_iter=conf["data_per_epoch"]) dev_dataset = ImageFolder( root=conf["data_dir"], transform=trans_dict["valid"], is_valid_file=lambda fn: fn.split(os.path.sep)[-2] in dev_lab) yield train_dataset, dev_dataset, train_sampler, conf["batch_size"] # MODEL, OPTIMIZER & LOSS def model_generator(): if conf["trunk_model"] in [ "ResNet-18", "ResNet-34", "ResNet-50", "ResNet-101", "ResNet-152" ]: depth = re.match("ResNet-(?P<depth>[0-9]+)", conf["trunk_model"]).group("depth") trunk = eval(f"models.resnet{depth}(pretrained=True)") trunk_output_size = trunk.fc.in_features trunk.fc = nn.Identity() elif conf["trunk_model"].startswith("timm:"): _model_name = conf["trunk_model"][5:] trunk = timm.create_model(_model_name, pretrained=True) trunk.reset_classifier(0) trunk_output_size = trunk.num_features p_dropout = param_gen.suggest_uniform("p_dropout", 0.0, 1.0) embedder = nn.Sequential(nn.Linear(trunk_output_size, conf["dim"]), nn.Dropout(p_dropout)) classifier = nn.Linear(conf["dim"], len(labels), bias=False) trunk.to("cuda") embedder.to("cuda") classifier.to("cuda") model_dict = { "trunk": trunk, "embedder": embedder, "classifier": classifier } #lr = param_gen.suggest_loguniform("model_lr", 1e-5, 1e-3) # beta1 = 0.9 # beta2 = 0.999 # eps = 1e-8 decay = param_gen.suggest_loguniform("model_decay", 1e-10, 1e-2) #lr = param_gen.suggest_loguniform("model_lr", 1e-5, 1e-2) lr = param_gen.suggest_loguniform("model_lr", 10**(-4.75), 10**(-3.75)) beta1 = 1. - param_gen.suggest_loguniform("model_beta1", 1e-3, 1.) beta2 = 1. - param_gen.suggest_loguniform("model_beta2", 1e-4, 1.) eps = param_gen.suggest_loguniform("model_eps", 1e-8, 1) optim_dict = { "trunk_optimizer": RAdam(trunk.parameters(), lr=lr, weight_decay=decay, betas=(beta1, beta2), eps=eps), "embedder_optimizer": RAdam(embedder.parameters(), lr=lr, weight_decay=decay, betas=(beta1, beta2), eps=eps), "classifier_optimizer": RAdam(classifier.parameters(), lr=lr, weight_decay=decay, betas=(beta1, beta2), eps=eps) } loss_type = conf["loss_type"] metric_loss_info = ALL_LOSSES[loss_type](param_gen, num_classes=len(labels), embedding_size=conf["dim"]) if "param" in metric_loss_info: # Add optimizer for loss lr = param_gen.suggest_loguniform("loss_lr", 1e-5, 1e-2) # beta1 = 1. - param_gen.suggest_loguniform("loss_beta1", 1e-3, 1.) # beta2 = 1. - param_gen.suggest_loguniform("loss_beta2", 1e-4, 1.) # eps = param_gen.suggest_loguniform("loss_eps", 1e-10, 1e-5) beta1 = 0.9 beta2 = 0.999 eps = 1e-8 opt = RAdam(metric_loss_info["loss"].parameters(), lr=lr, betas=(beta1, beta2), eps=eps) optim_dict["metric_loss_optimizer"] = opt loss_dict = { "metric_loss": metric_loss_info["loss"], "classifier_loss": torch.nn.CrossEntropyLoss() } #cls_weight = param_gen.suggest_loguniform("cls_weight", 0.01, 100.0) cls_weight = param_gen.suggest_loguniform("cls_weight", 0.1, 1.0) loss_weights = {"metric_loss": 1, "classifier_loss": cls_weight} return { "models": model_dict, "optimizers": optim_dict, "loss_funcs": loss_dict, "loss_weights": loss_weights } return { "modules": model_generator, "fold_generator": dataset_fold_generator }
def get_image_transform(mode, args, resize_size=256, crop_size=224): if mode == 'weak': if args.use_rrc_on_wa: if args.custom_scale: trfs = [ transforms.RandomResizedCrop(size=crop_size, scale=(0.2, 1.0)) ] else: trfs = [transforms.RandomResizedCrop(size=crop_size)] else: trfs = [ transforms.Resize((resize_size, resize_size)), transforms.RandomCrop(crop_size), ] trfs += [ transforms.RandomHorizontalFlip(), transforms.ToTensor(), ] elif mode == 'simclr': s = args.aug_strength color_jitter = transforms.ColorJitter(0.4 * s, 0.4 * s, 0.4 * s, 0.1 * s) if args.use_rrc: if args.custom_scale: trfs = [ transforms.RandomResizedCrop(size=crop_size, scale=(0.2, 1.0)) ] else: trfs = [transforms.RandomResizedCrop(size=crop_size)] else: trfs = [ transforms.Resize((resize_size, resize_size)), transforms.RandomCrop(crop_size), ] trfs += [ transforms.RandomHorizontalFlip(), ] if args.jitter: trfs.append(transforms.RandomApply([color_jitter], p=0.8)) if args.grayscale: trfs.append(transforms.RandomGrayscale(p=0.2)) if args.gaussblur: trfs.append(GaussianBlur(kernel_size=int(0.1 * crop_size))) trfs.append(transforms.ToTensor()) elif mode == 'randaug': trfs = [RandAugment(args.ra_n, args.ra_m)] if args.use_rrc: if args.custom_scale: trfs += [ transforms.RandomResizedCrop(size=crop_size, scale=(0.2, 1.0)) ] else: trfs += [transforms.RandomResizedCrop(size=crop_size)] else: trfs += [ transforms.Resize((resize_size, resize_size)), transforms.RandomCrop(crop_size), ] trfs += [ transforms.RandomHorizontalFlip(), transforms.ToTensor(), ] elif mode == 'test': trfs = [ transforms.Resize((resize_size, resize_size)), transforms.CenterCrop(crop_size), transforms.ToTensor() ] if args.norm_img: if args.norm_img_mode == 'whitening': normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) elif args.norm_img_mode == 'pmone': normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) trfs.append(normalize) return trfs
def get_dataset(dataset, batch_size, split=1.0, cutout_length=0, N=3, M=5, RandA=False): """ MEAN = [0.49139968, 0.48215827, 0.44653124] STD = [0.24703233, 0.24348505, 0.26158768] transf = [ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip() ] normalize = [ transforms.Resize(224), # 32 transforms.ToTensor(), transforms.Normalize(MEAN, STD) ] cutout = [] if cutout_length > 0: cutout.append(Cutout(cutout_length)) train_transform = transforms.Compose(transf + normalize + cutout) # if RandA: # Add RandAugment with N, M(hyperparameter) # train_transform.transforms.insert(0, RandAugment(N, M)) valid_transform = transforms.Compose(normalize) """ train_transform = transforms.Compose([ # transforms.Resize((224, 224)), # 224 Spop transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) if RandA: # Add RandAugment with N, M(hyperparameter) train_transform.transforms.insert(0, RandAugment(N, M)) valid_transform = transforms.Compose([ # transforms.Resize((224, 224)), # 224 Spos transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) if dataset == "cifar10": dataset_train = CIFAR10(root="./data", train=True, download=True, transform=train_transform) dataset_valid = CIFAR10(root="./data", train=False, download=True, transform=valid_transform) elif dataset == 'cifar100': dataset_train = CIFAR100(root="./data", train=True, download=True, transform=train_transform) dataset_valid = CIFAR100(root="./data", train=False, download=True, transform=valid_transform) else: raise NotImplementedError split_idx = 0 train_sampler = None if split < 1.0: sss = StratifiedShuffleSplit(n_splits=5, test_size=1 - split, random_state=0) sss = sss.split(list(range(len(dataset_train))), dataset_train.targets) for _ in range(split_idx + 1): train_idx, valid_idx = next(sss) train_sampler = SubsetRandomSampler(train_idx) valid_sampler = SubsetSampler(valid_idx) else: valid_sampler = SubsetSampler([]) train_loader = torch.utils.data.DataLoader( dataset_train, batch_size=batch_size, shuffle=True if train_sampler is None else False, num_workers=32, pin_memory=True, sampler=train_sampler, drop_last=True) # valid_loader = torch.utils.data.DataLoader( # dataset_valid, batch_size=batch_size, shuffle=True if train_sampler is None else False, num_workers=16, # pin_memory=True, sampler=valid_sampler, drop_last=True) valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True, drop_last=True) print('train/valid:{}/{},\t ' 'batchsize:{} * train-step/valid-step:{}/{} -> train/valid:{}/{},\t' 'Split train/valid:{}/{} '.format( len(dataset_train), len(dataset_valid), batch_size, int(len(train_loader)), int(len(valid_loader)), int(len(train_loader)) * batch_size, int(len(valid_loader)) * batch_size, len(train_loader) / (len(dataset_train) // batch_size), len(valid_loader) / (len(dataset_valid) // batch_size), )) train_dataprovider = DataIterator(train_loader) val_dataprovider = DataIterator(valid_loader) return train_dataprovider, val_dataprovider, int(len(train_loader)), int( len(valid_loader))
config = config(config_path) exp_name = config['exp_name'] print(f"EXP: {exp_name}") train_transforms = transforms.Compose([transforms.Resize((385, 505)), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) test_transforms = transforms.Compose([transforms.Resize((385, 505)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) train_transforms.transforms.insert(0, RandAugment(2, 14)) class2idx = config['class2idx'] batch = config['batch_size'] num_workers = config['num_workers'] train_data = datasets.ImageFolder(root = config['Train'], transform = train_transforms) valid_data = myset(config['Val'], config['val_txt'], class2idx, test_transforms) test_data = myset(config['Test'], config['test_txt'], class2idx, test_transforms) train_loader = torch.utils.data.DataLoader(train_data, batch_size = batch , num_workers = num_workers,shuffle = True) valid_loader = torch.utils.data.DataLoader(valid_data, batch_size = batch, num_workers = num_workers,shuffle = False) test_loader = torch.utils.data.DataLoader(test_data, batch_size = batch, num_workers = num_workers,shuffle = False) ms = [] model1 = res(3)
def __init__(self, img_size=(256, 256)): self.randaugment = torchvision.transforms.Compose([ RandAugment(10, 10), ]) print("Augment ")
from torchvision.transforms import transforms from RandAugment import RandAugment import visdom from PIL import Image viz = visdom.Visdom() transform_train = transforms.Compose([ transforms.RandomCrop(256), # transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) # Add RandAugment with N, M(hyperparameter) transform_train.transforms.insert(0, RandAugment(2, 10)) AB = Image.open("datasets/img2rain/trainA/norain-1.png").convert('RGB')
def cifar10_unsupervised_dataloaders(): print('Data Preparation') train_transform = Compose([ Pad(4), RandomCrop(32, fill=128), RandomHorizontalFlip(), ToTensor(), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), RandomErasing(scale=(0.1, 0.33)), ]) unsupervised_train_transformation = Compose([ Pad(4), RandomCrop(32, fill=128), ToTensor(), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) # RANDAUGMENT unsupervised_train_transformation.transforms.insert(0, RandAugment(3, 9)) test_transform = Compose([ ToTensor(), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) # Train dataset with and without labels cifar10_train_ds = datasets.CIFAR10('/data/', train=True, download=True) num_classes = len(cifar10_train_ds.classes) print('Loading dataset {0} for training -- Num_samples: {1}, num_classes: {2}'.format(datasets.CIFAR10.__name__,len(cifar10_train_ds),10)) labelled_indices = [] unlabelled_indices = [] indices = np.random.permutation(len(cifar10_train_ds)) class_counters = list([0] * num_classes) max_counter = 10000 // num_classes for i in indices: dp = cifar10_train_ds[i] if len(cifar10_train_ds) < sum(class_counters): unlabelled_indices.append(i) else: y = dp[1] c = class_counters[y] if c < max_counter: class_counters[y] += 1 labelled_indices.append(i) else: unlabelled_indices.append(i) # Labelled and unlabelled dataset train_labelled_ds = Subset(cifar10_train_ds, labelled_indices) train_labelled_ds_t = AddTransform(train_labelled_ds, train_transform) # unlabelled ds and aug ds train_unlabelled_ds = Subset(cifar10_train_ds, unlabelled_indices) train_unlabelled_ds = ConcatDataset([train_unlabelled_ds,train_labelled_ds]) # apply transformation for both train_unlabelled_ds_t = AddTransform(train_unlabelled_ds, train_transform) train_unlabelled_aug_ds_t = AddTransform(train_unlabelled_ds, unsupervised_train_transformation) print('Labelled dataset -- Num_samples: {0}, classes: {1}, \n Unsupervised dataset -- Num_samples {2}, Augmentation -- Num_samples: {3}' .format(len(train_labelled_ds_t), 10, len(train_unlabelled_ds_t), len(train_unlabelled_aug_ds_t))) # Data loader for labeled and unlabeled train dataset train_labelled = DataLoader( train_labelled_ds_t, batch_size=64, shuffle=False, num_workers=8, pin_memory=True ) train_unlabelled = DataLoader( train_unlabelled_ds_t, batch_size=64, shuffle=False, num_workers=8, pin_memory=True ) train_unlabelled_aug = DataLoader( train_unlabelled_aug_ds_t, batch_size=64, shuffle=False, num_workers=8, pin_memory=True ) # Data loader for test dataset cifar10_test_ds = datasets.CIFAR10('/data/', transform=test_transform, train=False, download=True) print('Test set -- Num_samples: {0}'.format(len(cifar10_test_ds))) test = DataLoader( cifar10_test_ds, batch_size=64, shuffle=False, num_workers=8, pin_memory=True ) return train_labelled, train_unlabelled, train_unlabelled_aug, test
def train(): data_transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), # ImageNetPolicy(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) data_transform.transforms.insert(0, RandAugment(13, 3)) train_set = IMAGE_Dataset(Path(DATASET_ROOT_train), data_transform) data_loader = DataLoader(dataset=train_set, batch_size=16, shuffle=True, num_workers=1) model = torch.hub.load('pytorch/vision:v0.6.0', 'googlenet', pretrained=True) # model = models.vgg19(pretrained=True) # final_in_features = model.classifier[6].in_features # model.classifier[6].out_features=train_set.num_classes model = model.cuda(CUDA_DEVICES) model.train() best_model_params = copy.deepcopy(model.state_dict()) best_acc = 0.0 num_epochs = 200 criterion = nn.CrossEntropyLoss() stepsize = 20 base_lr = 0.001 max_lr = 0.01 base_mm = 0.8 max_mm = 0.99 for epoch in range(num_epochs): # newlr = get_triangular_lr(epoch,stepsize,base_lr,max_lr) #mm=get_dynamic_momentum(epoch,stepsize,base_mm,max_mm) optimizer = torch.optim.SGD(params=model.parameters(), lr=0.001, momentum=0.9) print(f'Epoch: {epoch + 1}/{num_epochs}') print('-' * len(f'Epoch: {epoch + 1}/{num_epochs}')) training_loss = 0.0 training_corrects = 0 for i, (inputs, labels) in enumerate(data_loader): inputs = Variable(inputs.cuda(CUDA_DEVICES)) labels = Variable(labels.cuda(CUDA_DEVICES)) optimizer.zero_grad() outputs = model(inputs) _, preds = torch.max(outputs.data, 1) loss = criterion(outputs, labels) loss.backward() optimizer.step() training_loss += loss.item() * inputs.size(0) # print(training_loss) #revise loss.data[0]-->loss.item() training_corrects += torch.sum(preds == labels.data) #print(f'training_corrects: {training_corrects}') training_loss = training_loss / len(train_set) training_acc = training_corrects.double() / len(train_set) print( f'Training loss: {training_loss:.4f}\taccuracy: {training_acc:.4f}\n' ) test_acc = test(model) if test_acc > best_acc: best_acc = test_acc best_model_params = copy.deepcopy(model.state_dict()) model.load_state_dict(best_model_params) torch.save(model, f'model-{best_acc:.02f}-best_train_acc.pth')
def __init__(self, mdlParams, indSet, index=None): """ Args: mdlParams (dict): Configuration for loading indSet (string): Indicates train, val, test """ # Mdlparams self.mdlParams = mdlParams # Current indSet = 'trainInd'/'valInd'/'testInd' self.indSet = indSet if self.indSet == 'trainInd': self.root = mdlParams['dataDir'] + '/train' self.index = index elif self.indSet == 'valInd': self.root = mdlParams['dataDir'] + '/train' self.index = index else: self.root = mdlParams['dataDir'] + '/test' self.names_list = [] # Number of classes self.numClasses = mdlParams['numClasses'] # Size to crop self.crop_size = mdlParams['crop_size'] # Model input size self.input_size = (np.int32(mdlParams['input_size'][0]), np.int32(mdlParams['input_size'][1])) # Potential class balancing option self.balancing = mdlParams['balance_classes'] # Potential setMean and setStd self.setMean = mdlParams['setMean'].astype(np.float32) self.setStd = mdlParams['setStd'].astype(np.float32) # Only downsample self.only_downsmaple = mdlParams.get('only_downsmaple', True) # Meta csv self.meta_path = mdlParams['meta_path'] self.meta_df = pd.read_pickle(self.meta_path) self.subsets_size = [0, 0, 0, 0, 0, 0, 0] self.image_path = [] self.image_labels = [] self.image_meta = [] for img in os.listdir(self.root): id = img.split('.')[0] label = list(self.meta_df.loc[self.meta_df['image_id'] == id, 'dx'])[0] self.image_labels.append(CLASS_LABELS.index(label)) self.subsets_size[CLASS_LABELS.index(label)] += 1 self.image_meta.append( self.meta_df.loc[self.meta_df['image_id'] == id, self.mdlParams['meta_features']].to_numpy()) self.image_path.append(os.path.join(self.root, img)) self.image_path = np.asarray(self.image_path) self.image_labels = np.asarray(self.image_labels) self.image_meta = np.asarray(self.image_meta) if indSet == 'trainInd': if index is not None: self.image_path = self.image_path[index] self.image_labels = self.image_labels[index] self.image_meta = self.image_meta[index] all_transforms = [] if self.only_downsmaple: all_transforms.append(transforms.Resize(self.input_size)) else: all_transforms.append(transforms.Resize(self.crop_size)) all_transforms.append(transforms.RandomCrop(self.input_size)) if mdlParams.get('flip_lr_ud', False): all_transforms.append(transforms.RandomHorizontalFlip()) all_transforms.append(transforms.RandomVerticalFlip()) # Full rot if mdlParams.get('full_rot', 0) > 0: if mdlParams.get('scale', False): all_transforms.append( transforms.RandomChoice([ transforms.RandomAffine(mdlParams['full_rot'], scale=mdlParams['scale'], shear=mdlParams.get( 'shear', 0), resample=Image.NEAREST), transforms.RandomAffine(mdlParams['full_rot'], scale=mdlParams['scale'], shear=mdlParams.get( 'shear', 0), resample=Image.BICUBIC), transforms.RandomAffine(mdlParams['full_rot'], scale=mdlParams['scale'], shear=mdlParams.get( 'shear', 0), resample=Image.BILINEAR) ])) else: all_transforms.append( transforms.RandomChoice([ transforms.RandomRotation(mdlParams['full_rot'], resample=Image.NEAREST), transforms.RandomRotation(mdlParams['full_rot'], resample=Image.BICUBIC), transforms.RandomRotation(mdlParams['full_rot'], resample=Image.BILINEAR) ])) # Color distortion if mdlParams.get('full_color_distort') is not None: all_transforms.append( transforms.ColorJitter( brightness=mdlParams.get('brightness_aug', 32. / 255.), saturation=mdlParams.get('saturation_aug', 0.5), contrast=mdlParams.get('contrast_aug', 0.5), hue=mdlParams.get('hue_aug', 0.2))) else: all_transforms.append( transforms.ColorJitter(brightness=32. / 255., saturation=0.5)) # Autoaugment if self.mdlParams.get('randaugment', False): all_transforms.append( RandAugment(self.mdlParams.get('N'), self.mdlParams.get('M'))) # Cutout if self.mdlParams.get('cutout', 0) > 0: all_transforms.append( Cutout_v0(n_holes=1, length=self.mdlParams['cutout'])) # Normalize all_transforms.append(transforms.ToTensor()) all_transforms.append( transforms.Normalize(np.float32(self.mdlParams['setMean']), np.float32(self.mdlParams['setStd']))) # All transforms self.train_composed = transforms.Compose(all_transforms) self.valid_composed = transforms.Compose([ transforms.Resize(self.input_size), transforms.ToTensor(), transforms.Normalize( torch.from_numpy(self.setMean).float(), torch.from_numpy(self.setStd).float()) ]) self.image_path = duplist(self.image_path) self.image_labels = duplist(self.image_labels) self.image_meta = duplist(self.image_meta) else: if index is not None: self.image_path = self.image_path[index] self.image_labels = self.image_labels[index] self.image_meta = self.image_meta[index] self.valid_composed = transforms.Compose([ transforms.Resize(self.input_size), transforms.ToTensor(), transforms.Normalize( torch.from_numpy(self.setMean).float(), torch.from_numpy(self.setStd).float()) ])
def __init__(self, mdlParams, indSet, fold=-1): """ Args: mdlParams (dict): Configuration for loading indSet (string): Indicates train, val, test """ # Mdlparams self.mdlParams = mdlParams # Current indSet = 'trainInd'/'valInd'/'testInd' self.indSet = indSet if self.indSet == 'trainInd': self.root = mdlParams['dataDir'] + '/train' assert (fold > 0 and fold <= mdlParams['fold']) self.fold = fold elif self.indSet == 'valInd': self.root = mdlParams['dataDir'] + '/train' assert (fold > 0 and fold <= mdlParams['fold']) self.fold = fold else: self.root = mdlParams['dataDir'] + '/test' assert (fold < 0) self.names_list = [] # Number of classes self.numClasses = mdlParams['numClasses'] # Size to crop self.crop_size = mdlParams['crop_size'] # Model input size self.input_size = (np.int32(mdlParams['input_size'][0]), np.int32(mdlParams['input_size'][1])) # Potential class balancing option self.balancing = mdlParams['balance_classes'] # Potential setMean and setStd self.setMean = mdlParams['setMean'].astype(np.float32) self.setStd = mdlParams['setStd'].astype(np.float32) # Only downsample self.only_downsmaple = mdlParams.get('only_downsmaple', True) # Meta csv self.meta_path = mdlParams['meta_path'] self.meta_df = pd.read_pickle(self.meta_path) class_label = 0 self.subsets_size = [] self.image_path = [] for dir in os.listdir(self.root): subset_img_path = [] subset_name_list = [] dir_path = os.path.join(self.root, dir) for image in os.listdir(dir_path): subset_name_list.append({ 'id': image.split('.')[0], 'label': class_label, }) subset_img_path.append(os.path.join(dir_path)) subset_size = len(subset_img_path) self.names_list.append(subset_name_list) self.image_path.append(subset_img_path) self.subsets_size.append(subset_size) class_label += 1 if indSet == 'trainInd': new_names_list = [] new_image_path = [] for i in range(self.numClasses): print('Before folding: ' + str(self.subsets_size[i])) fold_len = self.subsets_size[i] // self.mdlParams['fold'] self.names_list[i] = self.names_list[i][:( self.fold - 1) * fold_len] + self.names_list[i][self.fold * fold_len:] self.image_path[i] = self.image_path[i][:( self.fold - 1) * fold_len] + self.image_path[i][self.fold * fold_len:] self.subsets_size[i] = len(self.names_list[i]) print('After folding: ' + str(self.subsets_size[i])) new_names_list += self.names_list[i] new_image_path += self.image_path[i] self.names_list = new_names_list self.image_path = new_image_path all_transforms = [] if self.only_downsmaple: all_transforms.append(transforms.Resize(self.input_size)) else: all_transforms.append(transforms.Resize(self.crop_size)) all_transforms.append(transforms.RandomCrop(self.input_size)) if mdlParams.get('flip_lr_ud', False): all_transforms.append(transforms.RandomHorizontalFlip()) all_transforms.append(transforms.RandomVerticalFlip()) # Full rot if mdlParams.get('full_rot', 0) > 0: if mdlParams.get('scale', False): all_transforms.append( transforms.RandomChoice([ transforms.RandomAffine(mdlParams['full_rot'], scale=mdlParams['scale'], shear=mdlParams.get( 'shear', 0), resample=Image.NEAREST), transforms.RandomAffine(mdlParams['full_rot'], scale=mdlParams['scale'], shear=mdlParams.get( 'shear', 0), resample=Image.BICUBIC), transforms.RandomAffine(mdlParams['full_rot'], scale=mdlParams['scale'], shear=mdlParams.get( 'shear', 0), resample=Image.BILINEAR) ])) else: all_transforms.append( transforms.RandomChoice([ transforms.RandomRotation(mdlParams['full_rot'], resample=Image.NEAREST), transforms.RandomRotation(mdlParams['full_rot'], resample=Image.BICUBIC), transforms.RandomRotation(mdlParams['full_rot'], resample=Image.BILINEAR) ])) # Color distortion if mdlParams.get('full_color_distort') is not None: all_transforms.append( transforms.ColorJitter( brightness=mdlParams.get('brightness_aug', 32. / 255.), saturation=mdlParams.get('saturation_aug', 0.5), contrast=mdlParams.get('contrast_aug', 0.5), hue=mdlParams.get('hue_aug', 0.2))) else: all_transforms.append( transforms.ColorJitter(brightness=32. / 255., saturation=0.5)) # Autoaugment if self.mdlParams.get('randaugment', False): all_transforms.append( RandAugment(self.mdlParams.get('N'), self.mdlParams.get('M'))) # Cutout if self.mdlParams.get('cutout', 0) > 0: all_transforms.append( Cutout_v0(n_holes=1, length=self.mdlParams['cutout'])) # Normalize all_transforms.append(transforms.ToTensor()) all_transforms.append( transforms.Normalize(np.float32(self.mdlParams['setMean']), np.float32(self.mdlParams['setStd']))) # All transforms self.composed = transforms.Compose(all_transforms) elif indSet == 'validInd': new_names_list = [] new_image_path = [] for i in range(self.numClasses): print('Before folding: ' + str(self.subsets_size[i])) fold_len = self.subsets_size[i] / self.fold self.names_list[i] = self.names_list[i][(self.fold - 1) * fold_len:self.fold * fold_len] self.image_path[i] = self.image_path[i][(self.fold - 1) * fold_len:self.fold * fold_len] self.subsets_size[i] = len(self.names_list[i]) print('After folding: ' + str(self.subsets_size[i])) new_names_list += self.names_list[i] new_image_path += self.image_path[i] self.names_list = new_names_list self.image_path = new_image_path self.composed = transforms.Compose([ transforms.Resize(self.input_size), transforms.ToTensor(), transforms.Normalize( torch.from_numpy(self.setMean).float(), torch.from_numpy(self.setStd).float()) ]) else: new_names_list = [] new_image_path = [] for i in range(self.numClasses): new_names_list += self.names_list[i] new_image_path += self.image_path[i] self.names_list = new_names_list self.image_path = new_image_path self.composed = transforms.Compose([ transforms.Resize(self.input_size), transforms.ToTensor(), transforms.Normalize( torch.from_numpy(self.setMean).float(), torch.from_numpy(self.setStd).float()) ])
from torchvision.transforms import transforms from RandAugment import RandAugment import cv2 import os from PIL import Image path = os.path.join(os.getcwd(), 'FMD', 'image', 'train') transform_train = transforms.Compose([ transforms.RandomHorizontalFlip(), # transforms.ToTensor(), # transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) # Add RandAugment with N, M(hyperparameter) transform_train.transforms.insert(0, RandAugment(1, 9)) for root, dirs, files in os.walk(path): for f in files: if f.endswith('.jpg'): for i in range(10): s = os.path.join(root, f) img = cv2.imread(s) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = Image.fromarray(img) img = transform_train(img) print(os.path.join(root, '{}_aug_{}'.format(i, f))) img.save(os.path.join(root, '{}_aug_{}'.format(i, f)))