scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=3) ################################################ # Dataset and train-test helpers ################################################ transform_val = transforms.Compose([ transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) data_root = args.IMAGENET_DIR train_ds = ImageNet(data_root, split='train', download=False, transform=transform_val) train_loader = torch.utils.data.DataLoader(train_ds, batch_size=args.BATCHSIZE, shuffle=True, drop_last=False, num_workers=args.NUM_WORKERS, pin_memory=True) val_ds = ImageNet(data_root, split='val', download=False, transform=transform_val) val_loader = torch.utils.data.DataLoader(val_ds, batch_size=args.BATCHSIZE, shuffle=True,
args = parse_args() # Location of test images images_fpaths = get_image_paths() if args.image_name is not None: idx = [ i for i in range(len(images_fpaths)) if args.image_name in images_fpaths[i] ][0] images_fpaths = [images_fpaths[idx]] # Get the actual labels from the filenames labels_act = parse_labels_from_paths(images_fpaths) # Get ImageNet dataset for interpreting the outputs of the network imagenet = ImageNet(IMAGENET_LOCATION, split='val', download=True) # For each input image, transform to be same size and orientation transform = init_image_transforms(args.image_size) # Put all images into a big tensor for putting into model imgs = read_images_as_tensors(images_fpaths, transform) # Instantiate model model = TransparentSqueezeNet(pretrained=True) # Pass images through model output_opaque = model(imgs) # Predicted label labels_pred = to_readable(output_opaque, imagenet)
dtype = args.dtype num_workers = args.num_workers batch_size = args.batch_size * len(device_ids) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]) if not args.use_lmdb: val_set = ImageNet(args.data_path, split='val', transform=val_transform) else: val_set = ImageLMDB(os.path.join(args.data_path, 'val.lmdb'), transform=val_transform) val_data = DataLoader(val_set, batch_size, False, pin_memory=True, num_workers=num_workers, drop_last=False) model_setting = set_model(0, args.norm_layer, args.activation) try: model = get_model(args.model, alpha=args.alpha, **model_setting)
def load_dataset(dataset_name, **kwargs): """ Loads the specified dataset and returns a PyTorch dataset object. Applies the standard transformations for said dataset by default. """ data_path = pathlib.Path('data').resolve() if dataset_name == 'cifar10': from torchvision.datasets import CIFAR10 # This is the standard normalization transformation transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) # User can specify to load the training set; loads the test set by default. train = kwargs.pop('train', False) dataset = CIFAR10(data_path, train=train, transform=transform, download=True) elif dataset_name == 'cifar100': from torchvision.datasets import CIFAR100 # This is the standard normalization transformation transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) # User can specify to load the training set; loads the test set by default. train = kwargs.pop('train', False) dataset = CIFAR100(data_path, train=train, transform=transform, download=True) elif dataset_name == 'imagenet': # Requires imagenet to be downloaded locally from torchvision.datasets import ImageNet # Standard transformation transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) dataset = ImageNet(data_path / 'imagenet', split='val', transform=transform) elif dataset_name == 'cifar10r': from data.nonstationary_datasets import CIFAR10R dataset = CIFAR10R() elif dataset_name == 'cifar100r': from data.nonstationary_datasets import CIFAR100R dataset = CIFAR100R() elif dataset_name == 'cifar10gb': from data.nonstationary_datasets import CIFAR10GB dataset = CIFAR10GB() elif dataset_name == 'cifar100gb': from data.nonstationary_datasets import CIFAR100GB dataset = CIFAR100GB() elif dataset_name == 'cifar10imba': from data.imbalanced_datasets import CIFAR10Imba dataset = CIFAR10Imba(class_ratios=kwargs['class_ratios']) else: raise NotImplementedError return dataset
def main(argv=None): path = '/home/alta/uncertainty/am969/pytorch/datasets/imagenet/tars' dataset1 = ImageNet(root=path, split='train', download=True) dataset2 = ImageNet(root=path, split='val ', download=True)
#input_transform = transforms.Compose([ # transforms.Resize(224,PIL.Image.BICUBIC), transforms.ToTensor(), # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) #]) input_transform = transforms.Compose([ transforms.Resize(256, PIL.Image.BICUBIC), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) test_dataset = ImageNet( DATA_ROOT, split="val", transform=input_transform, target_transform=None, download=True, ) test_loader = DataLoader( test_dataset, batch_size=64, shuffle=False, num_workers=4, pin_memory=True, ) model = model.cuda() model.eval()
def main_worker(gpu, ngpus_per_node, args): global best_acc1 args.gpu = gpu 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 if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) model = models.__dict__[args.arch](pretrained=True) else: print("=> creating model '{}'".format(args.arch)) model = getattr(resnet, args.arch)() if not torch.cuda.is_available(): print('using CPU, this will be slow') elif 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) else: # DataParallel will divide and allocate batch_size to all available GPUs if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() # 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'] best_acc1 = checkpoint['best_acc1'] if args.gpu is not None: # best_acc1 may be from a checkpoint from a different GPU best_acc1 = best_acc1.to(args.gpu) model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code img_net = ImageNet(args.data) traindir = os.path.join(args.data, 'train') valdir = os.path.join(args.data, 'val') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None 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) val_loader = torch.utils.data.DataLoader(datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: validate(val_loader, model, criterion, args) return for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch, args) # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args) # evaluate on validation set acc1 = validate(val_loader, model, criterion, args) # remember best acc@1 and save checkpoint is_best = acc1 > best_acc1 best_acc1 = max(acc1, best_acc1) if not args.multiprocessing_distributed or ( args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), }, is_best)
# -*- coding: utf-8 -*- # @Author : DevinYang([email protected]) import argparse from torchvision.datasets import ImageNet from torchtoolbox.tools import check_dir from torchtoolbox.tools.convert_lmdb import generate_lmdb_dataset, raw_reader parser = argparse.ArgumentParser(description='Convert a ImageFolder dataset to LMDB format.') parser.add_argument('--data-dir', type=str, required=True, help='ImageFolder path, this param will give to ImageFolder Dataset.') parser.add_argument('--save-dir', type=str, required=True, help='Save dir.') parser.add_argument('--download', action='store_true', help='download dataset.') parser.add_argument('-j', dest='num_workers', type=int, default=0) parser.add_argument('--write-frequency', type=int, default=5000) parser.add_argument('--max-size', type=float, default=1.0, help='Maximum size database, this is rate, default is 1T, final setting would be ' '1T * `this param`') args = parser.parse_args() check_dir(args.save_dir) train_data_set = ImageNet(args.data_dir, 'train', args.download, loader=raw_reader) val_data_set = ImageNet(args.data_dir, 'val', args.download, loader=raw_reader) if __name__ == '__main__': generate_lmdb_dataset(train_data_set, args.save_dir, 'train', args.num_workers, args.max_size, args.write_frequency) generate_lmdb_dataset(val_data_set, args.save_dir, 'val', args.num_workers, args.max_size, args.write_frequency)
def experiment(num_shared_classes, percent_shared_data, n_epochs=200,batch_size=128, eps=.3, adv_steps=100, learning_rate=.0004, gpu_num=1,adv_training=False,task="CIFAR100"): print("epochs,batch_size,eps,adv_steps,learning_rate,task") print(n_epochs,batch_size,eps,adv_steps,learning_rate,task) cuda = torch.cuda.is_available() transform_test = transforms.Compose( [transforms.ToTensor(),transforms.Normalize((0.5070751592371323, 0.48654887331495095, 0.4409178433670343), (0.2673342858792401, 0.2564384629170883, 0.27615047132568404))]) transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5070751592371323, 0.48654887331495095, 0.4409178433670343), (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)), ]) if task.upper() == "CIFAR100": transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) train_data = CIFAR100("data/",transform=transform_train, download=False) test_data = CIFAR100("data/", train=False, transform=transform_test, download=False) elif task.upper() == "IMAGENET": train_data = ImageNet('data/imagenet', split='train', download=False) test_data = ImageNet('data/imagenet', split='val', download=False) elif task.upper() == "FASHIONMNIST": transform = transforms.Compose([transforms.Lambda(lambda image: image.convert('RGB')), transforms.ToTensor() ]) train_data = FashionMNIST('data/fashionmnist',transform=transform, train=True, download=False) test_data = FashionMNIST('data/fashionmnist', transform=transform, train=False, download=False) else: train_data = CIFAR10("data/",transform=transform_train,download=False) test_data = CIFAR10("data/", train=False, transform=transform_test,download=False) # model1 = ResNet(ResidualBlock, [2, 2, 2],num_classes=10) # model2 = ResNet(ResidualBlock, [2, 2, 2],num_classes=10) all_classes = set([x[1] for x in train_data]) shared_classes = random.sample(all_classes, num_shared_classes) split_classes = [c for c in all_classes if c not in shared_classes] # get classes not shared if len(split_classes) % 2 == 1: # if we have an odd #, randomly remove one so that number of classes will be the same for each model split_classes.pop(random.randint(0, len(split_classes) - 1)) model1_split = random.sample(split_classes, len(split_classes) // 2) model2_split = [c for c in split_classes if c not in model1_split] model1_classes = model1_split model2_classes = model2_split model1_classes.sort() model2_classes.sort() # DEBUG: print("shared classes: {}".format(shared_classes)) print("model1 classes: {}".format(model1_classes)) print("model2 classes: {}".format(model2_classes)) model1_x_train = [] model1_y_train = [] model2_x_train = [] model2_y_train = [] shared_x_train = [] shared_y_train = [] # train data splits for index in range(len(train_data)): current_class = train_data[index][1] # model 1 if current_class in model1_classes: model1_x_train.append(train_data[index][0]) model1_y_train.append(train_data[index][1]) # model 2 if current_class in model2_classes: model2_x_train.append(train_data[index][0]) model2_y_train.append(train_data[index][1]) # split by percentage for classes per model1 if percent_shared_data < 100: new_model1_x_train = [] new_model1_y_train = [] for curr_class in model1_classes: temp_data_x = [] temp_data_y = [] # get all examples of class for i in range(len(model1_x_train)): if(model1_y_train[i] == curr_class): temp_data_x.append(model1_x_train[i]) temp_data_y.append(model1_y_train[i]) # split data by half the size total_size = len(temp_data_x) shared_size = int(total_size * .5) shared_indices = random.sample(list(range(len(temp_data_x))),shared_size) new_model1_x_train += [temp_data_x[i] for i in shared_indices] new_model1_y_train += [temp_data_y[i] for i in shared_indices] # split for model2 new_model2_x_train = [] new_model2_y_train = [] for curr_class in model2_classes: temp_data_x = [] temp_data_y = [] # get all examples of class for i in range(len(model2_x_train)): if(model2_y_train[i] == curr_class): temp_data_x.append(model2_x_train[i]) temp_data_y.append(model2_y_train[i]) # split data by half the size total_size = len(temp_data_x) shared_size = int(total_size * .5) shared_indices = random.sample(list(range(len(temp_data_x))),shared_size) new_model2_x_train += [temp_data_x[i] for i in shared_indices] new_model2_y_train += [temp_data_y[i] for i in shared_indices] # rewrite dataset model1_x_train = new_model1_x_train model1_y_train = new_model1_y_train model2_x_train = new_model2_x_train model2_y_train = new_model2_y_train # Carry out datasplitting for shared classes and add to datasets for shared_class in shared_classes: all_examples_x_train = [] all_examples_y_train = [] # get all examples of class for index in range(len(train_data)): current_class = train_data[index][1] if current_class == shared_class: all_examples_x_train.append(train_data[index][0]) all_examples_y_train.append(train_data[index][1]) # find max number of samples per model (set to be amount of examples if data is completely disjoint) max_examples = len(all_examples_x_train) // 2 # get shared examples shared_examples_x_train = [] shared_examples_y_train = [] num_shared_examples = max_examples * percent_shared_data // 100 for _ in range(num_shared_examples): random_int = random.randint(0, len(all_examples_x_train) - 1) shared_examples_x_train.append(all_examples_x_train.pop(random_int)) shared_examples_y_train.append(all_examples_y_train.pop(random_int)) # get disjoint examples disjoint_examples = max_examples - len(shared_examples_x_train) model1_examples_x_train = [] model1_examples_y_train = [] model2_examples_x_train = [] model2_examples_y_train = [] for _ in range(disjoint_examples): model1_rand_int = random.randint(0, len(all_examples_x_train) - 1) model1_examples_x_train.append(all_examples_x_train.pop(model1_rand_int)) model1_examples_y_train.append(all_examples_y_train.pop(model1_rand_int)) model2_rand_int = random.randint(0, len(all_examples_x_train) - 1) model2_examples_x_train.append(all_examples_x_train.pop(model2_rand_int)) model2_examples_y_train.append(all_examples_y_train.pop(model2_rand_int)) # add to the datasets for the model model1_x_train = shared_examples_x_train + model1_x_train + model1_examples_x_train model1_y_train = shared_examples_y_train + model1_y_train + model1_examples_y_train model2_x_train = shared_examples_x_train + model2_x_train + model2_examples_x_train model2_y_train = shared_examples_y_train + model2_y_train + model2_examples_y_train #print(model1_y_train) # assign mapping for new classes model1_class_mapping = {} model2_class_mapping = {} model1_classes_inc = 0 # go through model1 and assign unique classes to incrimental int starting at 0 for index in range(len(model1_y_train)): # if it doesn't exist assign if model1_y_train[index] not in model1_class_mapping.keys(): model1_class_mapping[model1_y_train[index]] = model1_classes_inc model1_classes_inc += 1 # append assigned token model1_y_train[index] = model1_class_mapping[model1_y_train[index]] model2_classes_inc = 0 # go through model2 and assign unique classes to incrimental int starting at 0 for index in range(len(model2_y_train)): # if it doesn't exist in model2 OR in model1, assign it if model2_y_train[index] not in model2_class_mapping.keys() and model2_y_train[index] not in model1_class_mapping.keys(): model2_class_mapping[model2_y_train[index]] = model2_classes_inc model2_y_train[index] = model2_classes_inc model2_classes_inc += 1 elif model2_y_train[index] in model1_class_mapping.keys(): model2_y_train[index] = model1_class_mapping[model2_y_train[index]] else: model2_y_train[index] = model2_class_mapping[model2_y_train[index]] model1_x_test = [] model1_y_test = [] model2_x_test = [] model2_y_test = [] shared_x_test = [] shared_y_test = [] # test data splits for index in range(len(test_data)): current_class = test_data[index][1] # model 1 if current_class in model1_classes: model1_x_test.append(test_data[index][0]) model1_y_test.append(test_data[index][1]) # model 2 if current_class in model2_classes: model2_x_test.append(test_data[index][0]) model2_y_test.append(test_data[index][1]) # shared classes for eval if current_class in shared_classes: shared_x_test.append(test_data[index][0]) shared_y_test.append(test_data[index][1]) model1_x_test += shared_x_test model1_y_test += shared_y_test model2_x_test += shared_x_test model2_y_test += shared_y_test for index in range(len(model1_y_test)): model1_y_test[index] = model1_class_mapping[model1_y_test[index]] for index in range(len(model2_y_test)): if model2_y_test[index] in model1_class_mapping.keys(): model2_y_test[index] = model1_class_mapping[model2_y_test[index]] else: model2_y_test[index] = model2_class_mapping[model2_y_test[index]] model1_classes_len= len(set([item for item in model1_y_train])) model2_classes_len = len(set([item for item in model2_y_train])) if task.upper() == "CIFAR100": model1 = models.wide_resnet50_2() model2 = models.wide_resnet50_2() # model1.fc = nn.Linear(2048, model1_classes_len) model2.fc = nn.Linear(2048, model2_classes_len) elif task.upper() == "IMAGENET": model1 = models.wide_resnet50_2() model2 = models.wide_resnet50_2() model1.fc = nn.Linear(2048, model1_classes_len) model2.fc = nn.Linear(2048, model2_classes_len) elif task.upper() == "FASHIONMNIST": model1 = models.resnet18() model2 = models.resnet18() model1.fc = nn.Linear(512, model1_classes_len) model2.fc = nn.Linear(512, model2_classes_len) else: # Get model (using ResNet50 for now) model1 = models.resnet50() model2 = models.resnet50() model1.fc = nn.Linear(2048, model1_classes_len) model2.fc = nn.Linear(2048, model2_classes_len) cuda = torch.cuda.is_available() if gpu_num in range(torch.cuda.device_count()): device = torch.device('cuda:'+str(gpu_num) if cuda else 'cpu') torch.cuda.set_device(device) else: device = torch.device('cpu') # Model Training model1 = model1.to(device) model2 = model2.to(device) criterion1 = nn.CrossEntropyLoss() optimizer1 = optim.AdamW(model1.parameters(), lr=learning_rate) scheduler1 = optim.lr_scheduler.MultiStepLR(optimizer1,milestones=[60, 120, 160], gamma=.2) #learning rate decay criterion2 = nn.CrossEntropyLoss() optimizer2 = optim.AdamW(model2.parameters(), lr=learning_rate) scheduler2 = optim.lr_scheduler.MultiStepLR(optimizer2,milestones=[60, 120, 160], gamma=.2) #learning rate decay # zip together two lists train_set1 = list(zip(model1_x_train, model1_y_train)) # create trainloader 1 trainloader_1 = torch.utils.data.DataLoader(train_set1, batch_size=batch_size, shuffle=True, num_workers=2) # create trainloader 2 # zip together two lists train_set2 = list(zip(model2_x_train, model2_y_train)) # create trainloader 1 trainloader_2 = torch.utils.data.DataLoader(train_set2, batch_size=batch_size, shuffle=True, num_workers=2) # TODO change this num_adv_batchs = 2 if adv_training else 0 adv_batches = random.sample(range(len(trainloader_1)), num_adv_batchs) #print("adv_batches:", adv_batches) # train model 1 for epoch in tqdm(range(n_epochs),desc="Epoch"): # loop over the dataset multiple times running_loss = 0.0 for i, data in enumerate(trainloader_1, 0): if cuda: data = tuple(d.cuda() for d in data) # get the inputs; data is a list of [inputs, labels] inputs, labels = data # zero the parameter gradients optimizer1.zero_grad() # forward + backward + optimize # train adversarial # if i in adv_batches: # print("adv training!") # adversary = LinfPGDAttack( # model1, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=eps, # nb_iter=adv_steps, eps_iter=0.01, rand_init=True, clip_min=0.0, clip_max=1.0, # targeted=False) # inputs = adversary.perturb(inputs, labels) outputs = model1(inputs) loss = criterion1(outputs, labels) loss.backward() optimizer1.step() # print statistics running_loss += loss.item() if i % 2000 == 1999: # print every 2000 mini-batches print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000)) running_loss = 0.0 print('Finished Training model1') # train model 2 for epoch in tqdm(range(n_epochs),desc="Epoch"): # loop over the dataset multiple times running_loss = 0.0 for i, data in enumerate(trainloader_2, 0): if cuda: data = tuple(d.cuda() for d in data) # get the inputs; data is a list of [inputs, labels] inputs, labels = data # zero the parameter gradients optimizer2.zero_grad() # forward + backward + optimize outputs = model2(inputs) loss = criterion2(outputs, labels) loss.backward() optimizer2.step() # print statistics running_loss += loss.item() if i % 2000 == 1999: # print every 2000 mini-batches print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000)) running_loss = 0.0 print('Finished Training model2') model1 = model1.to("cpu") model2 = model2.to("cpu") # convert shared classes to new labels for index in range(len(shared_y_test)): if shared_y_test[index] in model1_class_mapping.keys(): shared_y_test[index] = model1_class_mapping[shared_y_test[index]] else: shared_y_test[index] = model2_class_mapping[shared_y_test[index]] shared_y_test = torch.Tensor(shared_y_test).long() # if cuda: # shared_x_test = tuple(d.cuda() for d in shared_x_test) # shared_y_test = torch.Tensor(shared_y_test).long().cuda() model1_x_test = torch.stack(model1_x_test) model2_x_test = torch.stack(model2_x_test) model1.eval() shared_x_test = torch.stack(shared_x_test) model1.eval() adversary = LinfPGDAttack( model1, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=eps, nb_iter=adv_steps, eps_iter=0.01, rand_init=True, clip_min=0.0, clip_max=1.0, targeted=False) adv_untargeted = adversary.perturb(shared_x_test, shared_y_test) timestr = time.strftime("%Y%m%d_%H%M%S") print("saving models at", timestr) model1_name = './models/{}_{}_{}_model1_{}.pickle'.format(task,num_shared_classes, percent_shared_data,timestr) model2_name = './models/{}_{}_{}_model2_{}.pickle'.format(task,num_shared_classes, percent_shared_data,timestr) adv_name = './models/{}_{}_{}_adv_{}.pickle'.format(task,num_shared_classes, percent_shared_data,timestr) torch.save(model1, model1_name) torch.save(model2, model2_name) torch.save(adversary, adv_name) # Eval with torch.no_grad(): model1.eval() model2.eval() # model1 outputs output1 = model1(model1_x_test) shared_output1 = model1(shared_x_test) adv_output1 = model1(adv_untargeted) # model2 outputs output2 = model2(model2_x_test) shared_output2 = model2(shared_x_test) adv_output2 = model2(adv_untargeted) if task.upper() == "CIFAR100": # model 1 print("model1_acc:", accuracy(output1,model1_y_test)) print("model1_acc_5:", accuracy_n(output1,model1_y_test,5)) print("model1_acc_shared:", accuracy(shared_output1,shared_y_test)) print("model1_acc_5_shared:", accuracy_n(shared_output1,shared_y_test,5)) print("model1_adv_acc_shared:", accuracy(adv_output1,shared_y_test)) print("model1_adv_acc_5_shared:", accuracy_n(adv_output1,shared_y_test,5)) print() # model 2 print("model2_acc:", accuracy(output2,model2_y_test)) print("model2_acc_5:", accuracy_n(output2,model2_y_test,5)) print("model2_acc_shared:", accuracy(shared_output2,shared_y_test)) print("model2_acc_5_shared:", accuracy_n(shared_output2,shared_y_test,5)) print("model2_adv_acc_shared:", accuracy(adv_output2,shared_y_test)) print("model2_adv_acc_5_shared:", accuracy_n(adv_output2,shared_y_test,5)) else: # model 1 print("model1_acc:", accuracy(output1,model1_y_test)) print("model1_acc_shared:", accuracy(shared_output1,shared_y_test)) print("model1_adv_acc_shared:", accuracy(adv_output1,shared_y_test)) print() # model 2 print("model2_acc:", accuracy(output2,model2_y_test)) print("model2_acc_shared:", accuracy(shared_output2,shared_y_test)) print("model2_adv_acc_shared:", accuracy(adv_output2,shared_y_test))
def init_imagenet(dpath=IMAGENET_LOCATION): from torchvision.datasets import ImageNet imagenet = ImageNet(IMAGENET_LOCATION, split='val', download=True) return imagenet
def get_train_val_loaders( root_path: str, train_transforms: Callable, val_transforms: Callable, batch_size: int = 16, num_workers: int = 8, val_batch_size: Optional[int] = None, limit_train_num_samples: Optional[int] = None, limit_val_num_samples: Optional[int] = None, ) -> Tuple[DataLoader, DataLoader, DataLoader]: train_ds = ImageNet( root_path, split="train", transform=lambda sample: train_transforms(image=sample)["image"], loader=opencv_loader) val_ds = ImageNet( root_path, split="val", transform=lambda sample: val_transforms(image=sample)["image"], loader=opencv_loader) if limit_train_num_samples is not None: np.random.seed(limit_train_num_samples) train_indices = np.random.permutation( len(train_ds))[:limit_train_num_samples] train_ds = Subset(train_ds, train_indices) if limit_val_num_samples is not None: np.random.seed(limit_val_num_samples) val_indices = np.random.permutation( len(val_ds))[:limit_val_num_samples] val_ds = Subset(val_ds, val_indices) # random samples for evaluation on training dataset if len(val_ds) < len(train_ds): np.random.seed(len(val_ds)) train_eval_indices = np.random.permutation(len(train_ds))[:len(val_ds)] train_eval_ds = Subset(train_ds, train_eval_indices) else: train_eval_ds = train_ds train_loader = idist.auto_dataloader( train_ds, shuffle=True, batch_size=batch_size, num_workers=num_workers, drop_last=True, ) val_batch_size = batch_size * 4 if val_batch_size is None else val_batch_size val_loader = idist.auto_dataloader( val_ds, shuffle=False, batch_size=val_batch_size, num_workers=num_workers, drop_last=False, ) train_eval_loader = idist.auto_dataloader( train_eval_ds, shuffle=False, batch_size=val_batch_size, num_workers=num_workers, drop_last=False, ) return train_loader, val_loader, train_eval_loader
def main_worker(gpu, ngpus_per_node, args): args.gpu = gpu logger = get_logger(args.logging_file) logger.info("Use GPU: {} for training".format(args.gpu)) args.rank = args.rank * ngpus_per_node + gpu torch.distributed.init_process_group(backend="nccl", init_method=args.dist_url, world_size=args.world_size, rank=args.rank) epochs = args.epochs input_size = args.input_size resume_epoch = args.resume_epoch initializer = KaimingInitializer() zero_gamma = ZeroLastGamma() mix_precision_training = args.mix_precision_training is_first_rank = True if args.rank % ngpus_per_node == 0 else False batches_pre_epoch = args.num_training_samples // (args.batch_size * ngpus_per_node) lr = 0.1 * (args.batch_size * ngpus_per_node // 32) if args.lr == 0 else args.lr model = get_model(models, args.model) model.apply(initializer) if args.last_gamma: model.apply(zero_gamma) logger.info('Apply zero last gamma init.') if is_first_rank and args.model_info: summary(model, torch.rand((1, 3, input_size, input_size))) parameters = model.parameters() if not args.no_wd else no_decay_bias(model) if args.sgd_gc: logger.info('Use SGD_GC optimizer.') optimizer = SGD_GC(parameters, lr=lr, momentum=args.momentum, weight_decay=args.wd, nesterov=True) else: optimizer = optim.SGD(parameters, lr=lr, momentum=args.momentum, weight_decay=args.wd, nesterov=True) lr_scheduler = CosineWarmupLr(optimizer, batches_pre_epoch, epochs, base_lr=args.lr, warmup_epochs=args.warmup_epochs) # dropblock_scheduler = DropBlockScheduler(model, batches_pre_epoch, epochs) if args.lookahead: optimizer = Lookahead(optimizer) logger.info('Use lookahead optimizer.') torch.cuda.set_device(args.gpu) model.cuda(args.gpu) args.num_workers = int( (args.num_workers + ngpus_per_node - 1) / ngpus_per_node) if args.mix_precision_training and is_first_rank: logger.info('Train with FP16.') scaler = GradScaler(enabled=args.mix_precision_training) model = nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) Loss = nn.CrossEntropyLoss().cuda(args.gpu) if not args.label_smoothing else \ LabelSmoothingLoss(args.classes, smoothing=0.1).cuda(args.gpu) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if args.autoaugment: train_transform = transforms.Compose([ transforms.RandomResizedCrop(input_size), transforms.RandomHorizontalFlip(), ImageNetPolicy, transforms.ToTensor(), normalize, ]) else: train_transform = transforms.Compose([ transforms.RandomResizedCrop(input_size), # Cutout(), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.4, 0.4, 0.4), transforms.ToTensor(), normalize, ]) val_transform = transforms.Compose([ transforms.Resize(int(input_size / 0.875)), transforms.CenterCrop(input_size), transforms.ToTensor(), normalize, ]) train_set = ImageNet(args.data_path, split='train', transform=train_transform) val_set = ImageNet(args.data_path, split='val', transform=val_transform) train_sampler = DistributedSampler(train_set) train_loader = DataLoader(train_set, args.batch_size, False, pin_memory=True, num_workers=args.num_workers, drop_last=True, sampler=train_sampler) val_loader = DataLoader(val_set, args.batch_size, False, pin_memory=True, num_workers=args.num_workers, drop_last=False) if resume_epoch > 0: loc = 'cuda:{}'.format(args.gpu) checkpoint = torch.load(args.resume_param, map_location=loc) model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) scaler.load_state_dict(checkpoint['scaler']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) print("Finish loading resume param.") torch.backends.cudnn.benchmark = True top1_acc = metric.Accuracy(name='Top1 Accuracy') top5_acc = metric.TopKAccuracy(top=5, name='Top5 Accuracy') loss_record = metric.NumericalCost(name='Loss') for epoch in range(resume_epoch, epochs): tic = time.time() train_sampler.set_epoch(epoch) if not args.mixup: train_one_epoch(model, train_loader, Loss, optimizer, epoch, lr_scheduler, logger, top1_acc, loss_record, scaler, args) else: train_one_epoch_mixup(model, train_loader, Loss, optimizer, epoch, lr_scheduler, logger, loss_record, scaler, args) train_speed = int(args.num_training_samples // (time.time() - tic)) if is_first_rank: logger.info( 'Finish one epoch speed: {} samples/s'.format(train_speed)) test(model, val_loader, Loss, epoch, logger, top1_acc, top5_acc, loss_record, args) if args.rank % ngpus_per_node == 0: checkpoint = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scaler': scaler.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), } torch.save( checkpoint, '{}/{}_{}_{:.5}.pt'.format(args.save_dir, args.model, epoch, top1_acc.get()))
default=100, help='number of examples/minibatch') parser.add_argument('--num_batches', type=int, required=False, help='number of batches (default entire dataset)') args = parser.parse_args() model = resnet50(pretrained=True, progress=True) normalizer = utils.DifferentiableNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) dataset = ImageNet( args.imagenet_path, split='val', transform=transforms.Compose([ transforms.CenterCrop(224), transforms.ToTensor(), ]), ) val_loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, ) model.eval() if torch.cuda.is_available(): model.cuda() cw_loss = lf.CWLossF6(model, normalizer, kappa=float('inf')) perturbation_loss = lf.PerturbationNormLoss(lp=2)
def main(args): model = Model() optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) if args.scheduler=='multistep': scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, args.milestones, gamma=args.gamma) elif args.scheduler=='cosine': scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.step_ssize) criterion = torch.nn.CrossEntropyLoss() model = model.cuda() criterion = criterion.cuda() start_epoch = 0 # Check number of parameters your model pytorch_total_params = sum(p.numel() for p in model.parameters()) print(f"Number of parameters: {pytorch_total_params}") if not os.path.exists('{}'.format(args.savepath)): os.makedirs('{}'.format(args.savepath)) # resume if args.resume: model, optimizer, start_epoch = load_ckpt(model, optimizer, args) # Dataloader if args.dataset=='cifar10': normalize = transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) transform_val = transforms.Compose([ transforms.ToTensor(), normalize, ]) trainset = CIFAR10( root=args.datapath, train=True, download=True, transform=transform_train) valset = CIFAR10( root=args.datapath, train=False, download=True, transform=transform_val) elif args.dataset=='cifar100': normalize = transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) transform_val = transforms.Compose([ transforms.ToTensor(), normalize, ]) trainset = CIFAR100( root=args.datapath, train=True, download=True, transform=transform_train) valset = CIFAR100( root=args.datapath, train=False, download=True, transform=transform_val) elif args.dataset=='ImageNet': normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform_train = transforms.Compose([ transforms.RandomResizedCrop(image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) transform_val = transforms.Compose([ transforms.Resize(image_size + 32), transforms.CenterCrop(image_size), transforms.ToTensor(), normalize, ]) trainset = ImageNet( root=args.datapath, split='train', download=False, transform=transform_train) valset = ImageNet( root=args.datapath, split='val', download=False, transform=transform_val) elif args.dataeset=='tiny-imagenet-200': normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform_train = transforms.Compose([ transforms.RandomResizedCrop(image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) transform_val = transforms.Compose([ transforms.Resize(image_size + 32), transforms.CenterCrop(image_size), transforms.ToTensor(), normalize, ]) trainset = ImageFolder( root=args.datapath, split='train', download=False, transform=transform_train) valset = ImageFolder( root=args.datapath, split='val', download=False, transform=transform_val) train_loader = torch.utils.data.DataLoader( trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=False) val_loader = torch.utils.data.DataLoader( valset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False) # start training last_top1_acc = 0 acc1_valid = 0 best_acc1 = 0 is_best = False for epoch in range(start_epoch, args.epochs): print("\n----- epoch: {}, lr: {} -----".format( epoch, optimizer.param_groups[0]["lr"])) # train for one epoch start_time = time.time() last_top1_acc = train(train_loader, epoch, model, optimizer, criterion) elapsed_time = time.time() - start_time print('==> {:.2f} seconds to train this epoch\n'.format( elapsed_time)) # validate for one epoch start_time = time.time() acc1_valid = validate(val_loader, model, criterion) elapsed_time = time.time() - start_time print('==> {:.2f} seconds to validate this epoch\n'.format( elapsed_time)) # learning rate scheduling scheduler.step() is_best = acc1_valid > best_acc1 best_acc1 = max(acc1_valid, best_acc1) checkpoint = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() } save_ckpt(checkpoint, is_best, args) #if is_best: # torch.save(model.state_dict(), args.savepath+'model_weight_best.pth') # Save model each epoch #torch.save(model.state_dict(), args.savepath+'model_weight_epoch{}.pth'.format(epoch)) print(f"Last Top-1 Accuracy: {last_top1_acc}") print(f"Number of parameters: {pytorch_total_params}")
def __len__(self): return ImageNet.__len__(self)
else: optimizer = torch.optim.SGD(get_parameters(model), lr=0.05, momentum=momentum, nesterov=nesterov, weight_decay=4e-5) if args.train_mode == 'search': from nni.retiarii.oneshot.pytorch import ProxylessTrainer from torchvision.datasets import ImageNet normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) dataset = ImageNet(args.data_path, transform=transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) trainer = ProxylessTrainer(model, loss=LabelSmoothingLoss(), dataset=dataset, optimizer=optimizer, metrics=lambda output, target: accuracy( output, target, topk=( 1, 5, )), num_epochs=120, log_frequency=10) trainer.fit()
def create_dataset(name, root, split='validation', search_split=True, class_map=None, load_bytes=False, is_training=False, download=False, batch_size=None, repeats=0, **kwargs): """ Dataset factory method In parenthesis after each arg are the type of dataset supported for each arg, one of: * folder - default, timm folder (or tar) based ImageDataset * torch - torchvision based datasets * TFDS - Tensorflow-datasets wrapper in IterabeDataset interface via IterableImageDataset * all - any of the above Args: name: dataset name, empty is okay for folder based datasets root: root folder of dataset (all) split: dataset split (all) search_split: search for split specific child fold from root so one can specify `imagenet/` instead of `/imagenet/val`, etc on cmd line / config. (folder, torch/folder) class_map: specify class -> index mapping via text file or dict (folder) load_bytes: load data, return images as undecoded bytes (folder) download: download dataset if not present and supported (TFDS, torch) is_training: create dataset in train mode, this is different from the split. For Iterable / TDFS it enables shuffle, ignored for other datasets. (TFDS) batch_size: batch size hint for (TFDS) repeats: dataset repeats per iteration i.e. epoch (TFDS) **kwargs: other args to pass to dataset Returns: Dataset object """ name = name.lower() if name.startswith('torch/'): name = name.split('/', 2)[-1] torch_kwargs = dict(root=root, download=download, **kwargs) if name in _TORCH_BASIC_DS: ds_class = _TORCH_BASIC_DS[name] use_train = split in _TRAIN_SYNONYM ds = ds_class(train=use_train, **torch_kwargs) elif name == 'inaturalist' or name == 'inat': assert has_inaturalist, 'Please update to PyTorch 1.10, torchvision 0.11+ for Inaturalist' target_type = 'full' split_split = split.split('/') if len(split_split) > 1: target_type = split_split[0].split('_') if len(target_type) == 1: target_type = target_type[0] split = split_split[-1] if split in _TRAIN_SYNONYM: split = '2021_train' elif split in _EVAL_SYNONYM: split = '2021_valid' ds = INaturalist(version=split, target_type=target_type, **torch_kwargs) elif name == 'places365': assert has_places365, 'Please update to a newer PyTorch and torchvision for Places365 dataset.' if split in _TRAIN_SYNONYM: split = 'train-standard' elif split in _EVAL_SYNONYM: split = 'val' ds = Places365(split=split, **torch_kwargs) elif name == 'imagenet': if split in _EVAL_SYNONYM: split = 'val' ds = ImageNet(split=split, **torch_kwargs) elif name == 'image_folder' or name == 'folder': # in case torchvision ImageFolder is preferred over timm ImageDataset for some reason if search_split and os.path.isdir(root): # look for split specific sub-folder in root root = _search_split(root, split) ds = ImageFolder(root, **kwargs) else: assert False, f"Unknown torchvision dataset {name}" elif name.startswith('tfds/'): ds = IterableImageDataset(root, parser=name, split=split, is_training=is_training, download=download, batch_size=batch_size, repeats=repeats, **kwargs) else: # FIXME support more advance split cfg for ImageFolder/Tar datasets in the future if search_split and os.path.isdir(root): # look for split specific sub-folder in root root = _search_split(root, split) ds = ImageDataset(root, parser=name, class_map=class_map, load_bytes=load_bytes, **kwargs) return ds
def get_train_val_loaders( root_path: str, train_transforms: Callable, val_transforms: Callable, batch_size: int = 16, num_workers: int = 8, val_batch_size: Optional[int] = None, pin_memory: bool = True, random_seed: Optional[int] = None, train_sampler: Optional[Union[Sampler, str]] = None, val_sampler: Optional[Union[Sampler, str]] = None, limit_train_num_samples: Optional[int] = None, limit_val_num_samples: Optional[int] = None, ) -> Tuple[DataLoader, DataLoader, DataLoader]: train_ds = ImageNet( root_path, split="train", transform=lambda sample: train_transforms(image=sample)["image"], loader=opencv_loader) val_ds = ImageNet( root_path, split="val", transform=lambda sample: val_transforms(image=sample)["image"], loader=opencv_loader) if limit_train_num_samples is not None: if random_seed is not None: np.random.seed(random_seed) train_indices = np.random.permutation( len(train_ds))[:limit_train_num_samples] train_ds = Subset(train_ds, train_indices) if limit_val_num_samples is not None: val_indices = np.random.permutation( len(val_ds))[:limit_val_num_samples] val_ds = Subset(val_ds, val_indices) # random samples for evaluation on training dataset if len(val_ds) < len(train_ds): train_eval_indices = np.random.permutation(len(train_ds))[:len(val_ds)] train_eval_ds = Subset(train_ds, train_eval_indices) else: train_eval_ds = train_ds if isinstance(train_sampler, str): assert train_sampler == "distributed" train_sampler = data_dist.DistributedSampler(train_ds) train_eval_sampler = None if isinstance(val_sampler, str): assert val_sampler == "distributed" val_sampler = data_dist.DistributedSampler(val_ds, shuffle=False) train_eval_sampler = data_dist.DistributedSampler(train_eval_ds, shuffle=False) train_loader = DataLoader( train_ds, shuffle=train_sampler is None, batch_size=batch_size, num_workers=num_workers, sampler=train_sampler, pin_memory=pin_memory, drop_last=True, ) val_batch_size = batch_size * 4 if val_batch_size is None else val_batch_size val_loader = DataLoader( val_ds, shuffle=False, sampler=val_sampler, batch_size=val_batch_size, num_workers=num_workers, pin_memory=pin_memory, drop_last=False, ) train_eval_loader = DataLoader( train_eval_ds, shuffle=False, sampler=train_eval_sampler, batch_size=val_batch_size, num_workers=num_workers, pin_memory=pin_memory, drop_last=False, ) return train_loader, val_loader, train_eval_loader
import torch from torch.utils.data import DataLoader from torchvision.datasets import ImageNet print("Pytorch version:{}".format(torch.__version__)) #useful links #https://pytorch.org/docs/stable/_modules/torchvision/datasets/imagenet.html#ImageNet #https://pytorch.org/docs/stable/torchvision/datasets.html#imagenet #%% """ #_________________________________Dataloader for ImageNet______________________________# # # # The source code of the class ImageNet from pytorch says # # # # The dataset is no longer publicly accessible. # # You need to download the archives externally and place them in the root directory. # # # # I couldn't download the image because it takes 5 days to check a new account on # # the ImageNet website so i just implement the dataloader with the class ImageNet # # from pytorch # # # #______________________________________________________________________________________# """ root_dir = "path_to_ImageNet_directory_file" data = ImageNet(root=root_dir) data_loader = DataLoader(data, batch_size=1)
# Model model = vit_base_patch16_224(pretrained=True).cuda() baselines = Baselines(model) # LRP model_LRP = vit_LRP(pretrained=True).cuda() model_LRP.eval() lrp = LRP(model_LRP) # orig LRP model_orig_LRP = vit_orig_LRP(pretrained=True).cuda() model_orig_LRP.eval() orig_lrp = LRP(model_orig_LRP) # Dataset loader for sample images transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), ]) imagenet_ds = ImageNet(args.imagenet_validation_path, split='val', download=False, transform=transform) sample_loader = torch.utils.data.DataLoader( imagenet_ds, batch_size=args.batch_size, shuffle=False, num_workers=4 ) compute_saliency_and_save(args)
def _make_data_loader_imagenet(root, batch_size, workers=4, is_train=True, download=False, distributed=False): normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) logger = logging.getLogger('octconv') if is_train: logger.info("Loading ImageNet training data") st = time.time() scale = (0.08, 1.0) transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=scale), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) dataset = ImageNet(root=root, split='train', download=download, transform=transform) logger.info("Took: {}".format(time.time() - st)) if distributed: sampler = torch.utils.data.distributed.DistributedSampler(dataset) else: sampler = torch.utils.data.RandomSampler(dataset) loader = DataLoader(dataset, batch_size=batch_size, num_workers=workers, sampler=sampler, pin_memory=True) else: logger.info("Loading ImageNet validation data") transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]) dataset = ImageNet(root=root, split='val', download=download, transform=transform) if distributed: sampler = torch.utils.data.distributed.DistributedSampler(dataset) else: sampler = torch.utils.data.SequentialSampler(dataset) loader = DataLoader(dataset, batch_size=batch_size, num_workers=workers, sampler=sampler, pin_memory=True) return loader
args.batch_size = 1 model = vgg19(pretrained=True).to(device) args.BP = VanillaBackprop(model) elif args.method == 'integrad': args.batch_size = 1 model = vgg19(pretrained=True).to(device) args.BP = IntegratedGradients(model) elif args.method == 'gradcam': args.batch_size = 1 model = vgg19(pretrained=True).to(device) args.BP = GradCam(model, target_layer=36) # True for VGG-19 else: raise Exception('No model found') # Dataset loader for sample images transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor() ]) imagenet_ds = ImageNet('path/to/imagenet/', split='val', download=False, transform=transform) sample_loader = torch.utils.data.DataLoader( imagenet_ds, batch_size=args.batch_size, shuffle=False, num_workers=4 ) generate_visualization_to_hdf5(args)
def get_dataset(path: Path, transforms: Callable = None) -> Dataset: return ImageNet(str(path), split='val', transform=transforms)
def get_dataset(dataset='mnist', dataroot='data/', imresize=None, augment=False, normalize=False, tfm=None): num_chs = 1 if dataset.lower() in ['rsna', 'mnist'] else 3 valid_dataset = None if tfm is None: transform_lst = [] if imresize is not None: transform_lst.append(transforms.Resize(imresize)) transform_lst.append(transforms.ToTensor()) if normalize: transform_lst.append( transforms.Normalize(num_chs * (0.5, ), num_chs * (0.5, ))) transform = transforms.Compose(transform_lst) else: transform = tfm if dataset.lower() == 'rsna': raise NotImplementedError elif dataset.lower() == 'imagenet': train_dataset = ImageNet(dataroot, train=True, transform=transform, download=True) test_dataset = ImageNet(dataroot, train=False, transform=transform, download=True) elif dataset.lower() == 'cifar10': if tfm is None: transform = transform normalize_cifar10 = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if normalize: # overwrites the normalize transform = transforms.Compose([ transforms.ToTensor(), normalize_cifar10, ]) # else: # transform_train = transform train_dataset = CIFAR10(dataroot, train=True, transform=transform, download=True) test_dataset = CIFAR10(dataroot, train=False, transform=transform, download=True) elif dataset.lower() == 'cifar100': train_dataset = CIFAR100(dataroot, train=True, transform=transform, download=True) test_dataset = CIFAR100(dataroot, train=False, transform=transform, download=True) elif dataset.lower() == 'mnist': train_dataset = MNIST(dataroot, train=True, transform=transform, download=True) test_dataset = MNIST(dataroot, train=False, transform=transform, download=True) elif dataset.lower() == 'svhn': dataroot = os.path.join(dataroot, 'svhn') train_dataset = SVHN(dataroot, split='train', transform=transform, download=True) test_dataset = SVHN(dataroot, split='test', transform=transform, download=True) return train_dataset, test_dataset, num_chs