def main(): global args, best_prec1 args = parser.parse_args() args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size) # 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 = models.__dict__[args.arch](low_dim=args.low_dim) if not args.distributed: 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() else: model.cuda() model = torch.nn.parallel.DistributedDataParallel(model) # Data loading code 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.ImageFolderInstance( traindir, transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.2, 1.)), transforms.RandomGrayscale(p=0.2), transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), 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.ImageFolderInstance( 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) # define lemniscate and loss function (criterion) ndata = train_dataset.__len__() if args.nce_k > 0: lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m).cuda() criterion = NCECriterion(ndata).cuda() else: lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m).cuda() criterion = nn.CrossEntropyLoss().cuda() 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)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) lemniscate = checkpoint['lemniscate'] 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 if args.evaluate: kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t) return for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch) # train for one epoch train(train_loader, model, lemniscate, criterion, optimizer, epoch) # evaluate on validation set prec1 = NN(epoch, model, lemniscate, train_loader, val_loader) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'lemniscate': lemniscate, 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, is_best) # evaluate KNN after last epoch kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
def main(): # so other functions can access these variables global args, dataloaders, data_sizes, image_sets args = get_user_args() # defining processing device, if cuda is available then GPU else CPU device = torch.device("cuda:0" if ( torch.cuda.is_available() and args.gpu) else "cpu") print('=> beginning training using {}'.format(str(device).upper())) # lets the user know which model is being trained print('=> creating model: {}'.format(args.arch)) model = models.__dict__[args.arch](pretrained=True) print('* ' * 20) model.to(device) # send device to processor # image location with child folders of train, valid, test data_dir = Path(args.data) train_dir = data_dir / 'train' valid_dir = data_dir / 'valid' test_dir = data_dir / 'test' # variable for various iterations later states = ['train', 'valid', 'test'] # for easy iteration later dirs_dict = {'train': train_dir, 'valid': valid_dir, 'test': test_dir} # image normalization parameters, predefined normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # transforms for valid and test data, use same parameters valid_test_transforms = [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize] data_transforms = { 'train': # vector manipulation for generalized learning transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomRotation(30), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ToTensor(), normalize ]), 'valid': transforms.Compose(valid_test_transforms), 'test': transforms.Compose(valid_test_transforms) } image_sets = { i_set: datasets.ImageFolder( dirs_dict[i_set], transform=data_transforms[i_set]) for i_set in states } dataloaders = { 'train': torch.utils.data.DataLoader( image_sets['train'], batch_size=args.batch_size, shuffle=True), 'valid': torch.utils.data.DataLoader(image_sets['valid'], batch_size=args.batch_size), 'test': torch.utils.data.DataLoader(image_sets['test'], batch_size=args.batch_size) } classes = image_sets['train'].classes data_sizes = {x: len(image_sets[x]) for x in states} for p in model.parameters(): p.requires_grad = False # ensures gradients aren't calculated for parameters classifier = nn.Sequential( OrderedDict([ ('fc1', nn.Linear( model.classifier[0].in_features, args.hidden_units)), ('relu1,', nn.ReLU()), ('dropout', nn.Dropout(args.dropout)), ('fc2', nn.Linear(args.hidden_units, len(classes))), ('output', nn.LogSoftmax(dim=1)), ])) model.classifier = classifier criterion = nn.NLLLoss() optimizer = optim.Adam(model.classifier.parameters(), lr=args.learning_rate) scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.0125) model_trained = train(model, optimizer, criterion, scheduler, args.epochs, device) save_checkpoint(model_trained, args.epochs, args.save_dir, args.arch, args.learning_rate, optimizer, args.hidden_units)
# import matplotlib.pyplot as plt from tqdm import tqdm # Displays a progress bar import torch from torch import nn from torch import optim import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import Dataset, Subset, DataLoader, random_split # Load the dataset and train, val, test splits # TODO: currently only experiment with FashionMNIST. Need to process your own dataset print("Loading datasets...") DATA_transform= transforms.Compose([ transforms.RandomResizedCrop(224), transforms.ToTensor(), # Transform from [0,255] uint8 to [0,1] float transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize to zero mean and unit variance ]) training_data_dir='./data/Training' testing_data_dir='./data/Testing' TUMOR_train=datasets.ImageFolder(training_data_dir, DATA_transform) TUMOR_test = datasets.ImageFolder(testing_data_dir, DATA_transform) print(TUMOR_test.classes) # FASHION_trainval = datasets.FashionMNIST('.', download=True, train=True, transform=FASHION_transform) # FASHION_train = Subset(FASHION_trainval, range(50000)) # FASHION_val = Subset(FASHION_trainval, range(50000,60000)) # FASHION_test = datasets.FashionMNIST('.', download=True, train=False, transform=FASHION_transform) print("Done!")
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) if args.gen_map: args.qw = -1 args.qa = -1 # create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) else: print("=> creating model '{}'".format(args.arch)) try: model = models.__dict__[args.arch](pretrained=args.pretrained) args.qw = -1 args.qa = -1 except KeyError: if 'bp' in args.arch: model = imagenet_extra_models.__dict__[args.arch]( pretrained=args.pretrained, nbits_w=args.qw, log=args.quan_log, increase_factor=args.increase_factor) else: model = imagenet_extra_models.__dict__[args.arch]( pretrained=args.pretrained) print('model:\n=========\n{}\n=========='.format(model)) if args.gen_map: try: ori_model = models.__dict__[args.original_model]() except KeyError: ori_model = imagenet_extra_models.__dict__[args.original_model]() print('original model:\n=========\n{}\n=========='.format(ori_model)) key_map = gen_key_map(model.state_dict(), ori_model.state_dict()) with open('models/imagenet/{}_map.json'.format(args.arch), 'w') as wf: json.dump(key_map, wf) print('Generate key map done') return 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) 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) params = add_weight_decay(model, weight_decay=args.weight_decay, skip_keys=['expand_']) optimizer = torch.optim.SGD(params, args.lr, momentum=args.momentum) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume, map_location='cpu') model.load_state_dict(checkpoint['state_dict']) # if not args.quant_bias_scale: # 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) # try: # model.load_state_dict(checkpoint['state_dict']) # # ValueError: loaded state dict has a different number of parameter groups # # different version # # optimizer.load_state_dict(checkpoint['optimizer']) # except RuntimeError: # print('Fine-tune qfi_wide model using qfn_relaxed weights.') # key_map = gen_key_map(model.state_dict(), checkpoint['state_dict']) # load_fake_quantized_state_dict(model, checkpoint['state_dict'], key_map) # args.start_epoch = 0 # best_acc1 = 0 # optimizer = torch.optim.SGD(params, args.lr, # momentum=args.momentum) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code 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 args.log_name = 'logger/{}_{}'.format(args.arch, args.log_name) writer = SummaryWriter(args.log_name) with open('{}/{}.txt'.format(args.log_name, args.arch), 'w') as wf: wf.write(str(model)) scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs) scheduler_warmup = GradualWarmupScheduler(optimizer, multiplier=10, total_epoch=3, after_scheduler=scheduler_cosine) for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) # adjust_learning_rate(optimizer, epoch, args) scheduler_warmup.step() # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args, writer) # evaluate on validation set acc1, acc5 = validate(val_loader, model, criterion, args) writer.add_scalar('val/acc1', acc1, epoch) writer.add_scalar('val/acc5', acc5, epoch) # writer.add_scalar('val/bs', bs, epoch) # 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, prefix='{}/{}'.format(args.log_name, args.arch)) if epoch % 10 == 0: save_checkpoint_backup(model.state_dict(), prefix='{}/{}_{}'.format( args.log_name, args.arch, epoch)) writer.close()
def main(): # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example') parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', help='input batch size for testing (default: 1000)') parser.add_argument('--world_size', type=int, default=1, help='number of GPUs to use') parser.add_argument('--epochs', type=int, default=10, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--wd', type=float, default=1e-4, help='weight decay (default: 5e-4)') parser.add_argument('--lr-decay-every', type=int, default=100, help='learning rate decay by 10 every X epochs') parser.add_argument('--lr-decay-scalar', type=float, default=0.1, help='--') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)') parser.add_argument( '--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--run_test', default=False, type=str2bool, nargs='?', help='run test only') parser.add_argument( '--limit_training_batches', type=int, default=-1, help='how many batches to do per training, -1 means as many as possible' ) parser.add_argument('--no_grad_clip', default=False, type=str2bool, nargs='?', help='turn off gradient clipping') parser.add_argument('--get_flops', default=False, type=str2bool, nargs='?', help='add hooks to compute flops') parser.add_argument( '--get_inference_time', default=False, type=str2bool, nargs='?', help='runs valid multiple times and reports the result') parser.add_argument('--mgpu', default=False, type=str2bool, nargs='?', help='use data paralization via multiple GPUs') parser.add_argument('--dataset', default="MNIST", type=str, help='dataset for experiment, choice: MNIST, CIFAR10') parser.add_argument('--data', metavar='DIR', default='/imagenet', help='path to imagenet dataset') parser.add_argument( '--model', default="lenet3", type=str, help='model selection, choices: lenet3, vgg, mobilenetv2, resnet18', choices=[ "lenet3", "vgg", "mobilenetv2", "resnet18", "resnet152", "resnet50", "resnet50_noskip", "resnet20", "resnet34", "resnet101", "resnet101_noskip", "densenet201_imagenet", 'densenet121_imagenet', "multprun_gate5_gpu_0316_1", "mult_prun8_gpu", "multnas5_gpu" ]) parser.add_argument('--tensorboard', type=str2bool, nargs='?', help='Log progress to TensorBoard') parser.add_argument( '--save_models', default=True, type=str2bool, nargs='?', help='if True, models will be saved to the local folder') parser.add_argument('--fineturn_model', type=str2bool, nargs='?', help='Log progress to TensorBoard') # ============================PRUNING added parser.add_argument( '--pruning_config', default=None, type=str, help= 'path to pruning configuration file, will overwrite all pruning parameters in arguments' ) parser.add_argument('--group_wd_coeff', type=float, default=0.0, help='group weight decay') parser.add_argument('--name', default='test', type=str, help='experiment name(folder) to store logs') parser.add_argument( '--augment', default=False, type=str2bool, nargs='?', help= 'enable or not augmentation of training dataset, only for CIFAR, def False' ) parser.add_argument('--load_model', default='', type=str, help='path to model weights') parser.add_argument('--pruning', default=False, type=str2bool, nargs='?', help='enable or not pruning, def False') parser.add_argument( '--pruning-threshold', '--pt', default=100.0, type=float, help= 'Max error perc on validation set while pruning (default: 100.0 means always prune)' ) parser.add_argument( '--pruning-momentum', default=0.0, type=float, help= 'Use momentum on criteria between pruning iterations, def 0.0 means no momentum' ) parser.add_argument('--pruning-step', default=15, type=int, help='How often to check loss and do pruning step') parser.add_argument('--prune_per_iteration', default=10, type=int, help='How many neurons to remove at each iteration') parser.add_argument( '--fixed_layer', default=-1, type=int, help='Prune only a given layer with index, use -1 to prune all') parser.add_argument('--start_pruning_after_n_iterations', default=0, type=int, help='from which iteration to start pruning') parser.add_argument('--maximum_pruning_iterations', default=1e8, type=int, help='maximum pruning iterations') parser.add_argument('--starting_neuron', default=0, type=int, help='starting position for oracle pruning') parser.add_argument('--prune_neurons_max', default=-1, type=int, help='prune_neurons_max') parser.add_argument('--pruning-method', default=0, type=int, help='pruning method to be used, see readme.md') parser.add_argument('--pruning_fixed_criteria', default=False, type=str2bool, nargs='?', help='enable or not criteria reevaluation, def False') parser.add_argument('--fixed_network', default=False, type=str2bool, nargs='?', help='fix network for oracle or criteria computation') parser.add_argument( '--zero_lr_for_epochs', default=-1, type=int, help='Learning rate will be set to 0 for given number of updates') parser.add_argument( '--dynamic_network', default=False, type=str2bool, nargs='?', help= 'Creates a new network graph from pruned model, works with ResNet-101 only' ) parser.add_argument('--use_test_as_train', default=False, type=str2bool, nargs='?', help='use testing dataset instead of training') parser.add_argument('--pruning_mask_from', default='', type=str, help='path to mask file precomputed') parser.add_argument( '--compute_flops', default=True, type=str2bool, nargs='?', help= 'if True, will run dummy inference of batch 1 before training to get conv sizes' ) # ============================END pruning added best_prec1 = 0 global global_iteration global group_wd_optimizer global_iteration = 0 args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() torch.manual_seed(args.seed) args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=0) device = torch.device("cuda" if use_cuda else "cpu") # dataset loading section if args.dataset == "MNIST": kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} train_loader = torch.utils.data.DataLoader(datasets.MNIST( '../data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ])), batch_size=args.batch_size, shuffle=True, **kwargs) test_loader = torch.utils.data.DataLoader( datasets.MNIST('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ])), batch_size=args.test_batch_size, shuffle=True, **kwargs) elif args.dataset == "CIFAR10": # Data loading code normalize = transforms.Normalize( mean=[x / 255.0 for x in [125.3, 123.0, 113.9]], std=[x / 255.0 for x in [63.0, 62.1, 66.7]]) if args.augment: transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) else: transform_train = transforms.Compose([ transforms.ToTensor(), normalize, ]) transform_test = transforms.Compose([transforms.ToTensor(), normalize]) kwargs = {'num_workers': 8, 'pin_memory': True} train_loader = torch.utils.data.DataLoader(datasets.CIFAR10( '../data', train=True, download=True, transform=transform_train), batch_size=args.batch_size, shuffle=True, drop_last=True, **kwargs) test_loader = torch.utils.data.DataLoader( datasets.CIFAR10('../data', train=False, transform=transform_test), batch_size=args.test_batch_size, shuffle=True, **kwargs) elif args.dataset == "Imagenet": 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 kwargs = {'num_workers': 16} train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, pin_memory=True, **kwargs) if args.use_test_as_train: train_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=(train_sampler is None), **kwargs) test_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, pin_memory=True, **kwargs) #wm elif args.dataset == "mult_5T": args.data_root = ['/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/CX_20200709',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TK_20200709',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/ZR_20200709',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TX_20200616',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/WM_20200709'] args.data_root_val = ['/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/CX_20200709',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TK_20200709',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/ZR_20200709',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TX_20200616',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/WM_20200709'] args.train_data_list = ['/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/CX_20200709/txt/cx_train.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TK_20200709/txt/tk_train.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/ZR_20200709/txt/zr_train.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TX_20200616/txt/tx_train.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/WM_20200709/txt/wm_train.txt'] args.val_data_list = ['/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/CX_20200709/txt/cx_val.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TK_20200709/txt/tk_val.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/ZR_20200709/txt/zr_val.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/TX_20200616/txt/tx_val.txt',\ '/workspace/mnt/storage/yangdecheng/yangdecheng/data/TR-NMA-07/WM_20200709/txt/wm_val.txt'] num_tasks = len(args.data_root) args.ngpu = 8 args.workers = 8 args.train_batch_size = [40, 40, 40, 40, 40] #36 args.val_batch_size = [100, 100, 100, 100, 100] args.loss_weight = [1.0, 1.0, 1.0, 1.0, 1.0] args.val_num_classes = [[0, 1, 2, 3, 4], [0, 1, 2], [0, 1], [0, 1], [0, 1, 2, 3, 4, 5, 6]] args.mixup_alpha = None #None for i in range(num_tasks): args.train_batch_size[i] *= args.ngpu args.val_batch_size[i] *= args.ngpu pixel_mean = [0.406, 0.456, 0.485] pixel_std = [0.225, 0.224, 0.229] #私人定制: train_dataset = [] for i in range(num_tasks): if i == 1: train_dataset.append( FileListLabeledDataset( args.train_data_list[i], args.data_root[i], Compose([ RandomResizedCrop( 112, scale=(0.94, 1.), ratio=(1. / 4., 4. / 1.) ), #scale=(0.7, 1.2), ratio=(1. / 1., 4. / 1.) RandomHorizontalFlip(), ColorJitter(brightness=[0.5, 1.5], contrast=[0.5, 1.5], saturation=[0.5, 1.5], hue=0), ToTensor(), Lighting(1, [0.2175, 0.0188, 0.0045], [[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140], [-0.5836, -0.6948, 0.4203]]), Normalize(pixel_mean, pixel_std), ]))) else: train_dataset.append( FileListLabeledDataset( args.train_data_list[i], args.data_root[i], Compose([ RandomResizedCrop(112, scale=(0.7, 1.2), ratio=(1. / 1., 4. / 1.)), RandomHorizontalFlip(), ColorJitter(brightness=[0.5, 1.5], contrast=[0.5, 1.5], saturation=[0.5, 1.5], hue=0), ToTensor(), Lighting(1, [0.2175, 0.0188, 0.0045], [[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140], [-0.5836, -0.6948, 0.4203]]), Normalize(pixel_mean, pixel_std), ]))) #原来的 # train_dataset = [FileListLabeledDataset( # args.train_data_list[i], args.data_root[i], # Compose([ # RandomResizedCrop(112,scale=(0.7, 1.2), ratio=(1. / 1., 4. / 1.)), # RandomHorizontalFlip(), # ColorJitter(brightness=[0.5,1.5], contrast=[0.5,1.5], saturation=[0.5,1.5], hue= 0), # ToTensor(), # Lighting(1, [0.2175, 0.0188, 0.0045], [[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140], [-0.5836, -0.6948, 0.4203]]), # Normalize(pixel_mean, pixel_std),]), # memcached=False, # memcached_client="") for i in range(num_tasks)] args.num_classes = [td.num_class for td in train_dataset] train_longest_size = max([ int(np.ceil(len(td) / float(bs))) for td, bs in zip(train_dataset, args.train_batch_size) ]) train_sampler = [ GivenSizeSampler(td, total_size=train_longest_size * bs, rand_seed=0) for td, bs in zip(train_dataset, args.train_batch_size) ] train_loader = [ DataLoader(train_dataset[k], batch_size=args.train_batch_size[k], shuffle=False, num_workers=args.workers, pin_memory=False, sampler=train_sampler[k]) for k in range(num_tasks) ] val_dataset = [ FileListLabeledDataset( args.val_data_list[i], args.data_root_val[i], Compose([ Resize((112, 112)), # CenterCrop(112), ToTensor(), Normalize(pixel_mean, pixel_std), ]), memcached=False, memcached_client="") for i in range(num_tasks) ] val_longest_size = max([ int(np.ceil(len(vd) / float(bs))) for vd, bs in zip(val_dataset, args.val_batch_size) ]) test_loader = [ DataLoader(val_dataset[k], batch_size=args.val_batch_size[k], shuffle=False, num_workers=args.workers, pin_memory=False) for k in range(num_tasks) ] if args.model == "lenet3": model = LeNet(dataset=args.dataset) elif args.model == "vgg": model = vgg11_bn(pretrained=True) elif args.model == "resnet18": model = PreActResNet18() elif (args.model == "resnet50") or (args.model == "resnet50_noskip"): if args.dataset == "CIFAR10": model = PreActResNet50(dataset=args.dataset) else: from models.resnet import resnet50 skip_gate = True if "noskip" in args.model: skip_gate = False if args.pruning_method not in [22, 40]: skip_gate = False model = resnet50(skip_gate=skip_gate) elif args.model == "resnet34": if not (args.dataset == "CIFAR10"): from models.resnet import resnet34 model = resnet34() elif args.model == "multprun_gate5_gpu_0316_1": from models.multitask import MultiTaskWithLoss model = MultiTaskWithLoss(backbone=args.model, num_classes=args.num_classes, feature_dim=2560, spatial_size=112, arc_fc=False, feat_bn=False) print(model) elif args.model == "mult_prun8_gpu": from models.multitask import MultiTaskWithLoss model = MultiTaskWithLoss(backbone=args.model, num_classes=args.num_classes, feature_dim=18, spatial_size=112, arc_fc=False, feat_bn=False) print(model) elif args.model == "multnas5_gpu": #作为修改项 from models.multitask import MultiTaskWithLoss model = MultiTaskWithLoss(backbone=args.model, num_classes=args.num_classes, feature_dim=512, spatial_size=112, arc_fc=False, feat_bn=False) print(model) elif "resnet101" in args.model: if not (args.dataset == "CIFAR10"): from models.resnet import resnet101 if args.dataset == "Imagenet": classes = 1000 if "noskip" in args.model: model = resnet101(num_classes=classes, skip_gate=False) else: model = resnet101(num_classes=classes) elif args.model == "resnet20": if args.dataset == "CIFAR10": NotImplementedError( "resnet20 is not implemented in the current project") # from models.resnet_cifar import resnet20 # model = resnet20() elif args.model == "resnet152": model = PreActResNet152() elif args.model == "densenet201_imagenet": from models.densenet_imagenet import DenseNet201 model = DenseNet201(gate_types=['output_bn'], pretrained=True) elif args.model == "densenet121_imagenet": from models.densenet_imagenet import DenseNet121 model = DenseNet121(gate_types=['output_bn'], pretrained=True) else: print(args.model, "model is not supported") ####end dataset preparation if args.dynamic_network: # attempts to load pruned model and modify it be removing pruned channels # works for resnet101 only if (len(args.load_model) > 0) and (args.dynamic_network): if os.path.isfile(args.load_model): load_model_pytorch(model, args.load_model, args.model) else: print("=> no checkpoint found at '{}'".format(args.load_model)) exit() dynamic_network_change_local(model) # save the model log_save_folder = "%s" % args.name if not os.path.exists(log_save_folder): os.makedirs(log_save_folder) if not os.path.exists("%s/models" % (log_save_folder)): os.makedirs("%s/models" % (log_save_folder)) model_save_path = "%s/models/pruned.weights" % (log_save_folder) model_state_dict = model.state_dict() if args.save_models: save_checkpoint({'state_dict': model_state_dict}, False, filename=model_save_path) print("model is defined") # aux function to get size of feature maps # First it adds hooks for each conv layer # Then runs inference with 1 image output_sizes = get_conv_sizes(args, model) if use_cuda and not args.mgpu: model = model.to(device) elif args.distributed: model.cuda() print( "\n\n WARNING: distributed pruning was not verified and might not work correctly" ) model = torch.nn.parallel.DistributedDataParallel(model) elif args.mgpu: model = torch.nn.DataParallel(model).cuda() else: model = model.to(device) print( "model is set to device: use_cuda {}, args.mgpu {}, agrs.distributed {}" .format(use_cuda, args.mgpu, args.distributed)) weight_decay = args.wd if args.fixed_network: weight_decay = 0.0 # remove updates from gate layers, because we want them to be 0 or 1 constantly if 1: parameters_for_update = [] parameters_for_update_named = [] for name, m in model.named_parameters(): if "gate" not in name: parameters_for_update.append(m) parameters_for_update_named.append((name, m)) else: print("skipping parameter", name, "shape:", m.shape) total_size_params = sum( [np.prod(par.shape) for par in parameters_for_update]) print("Total number of parameters, w/o usage of bn consts: ", total_size_params) optimizer = optim.SGD(parameters_for_update, lr=args.lr, momentum=args.momentum, weight_decay=weight_decay) if 1: # helping optimizer to implement group lasso (with very small weight that doesn't affect training) # will be used to calculate number of remaining flops and parameters in the network group_wd_optimizer = group_lasso_decay( parameters_for_update, group_lasso_weight=args.group_wd_coeff, named_parameters=parameters_for_update_named, output_sizes=output_sizes) cudnn.benchmark = True # define objective criterion = nn.CrossEntropyLoss() ###=======================added for pruning # logging part log_save_folder = "%s" % args.name if not os.path.exists(log_save_folder): os.makedirs(log_save_folder) if not os.path.exists("%s/models" % (log_save_folder)): os.makedirs("%s/models" % (log_save_folder)) train_writer = None if args.tensorboard: try: # tensorboardX v1.6 train_writer = SummaryWriter(log_dir="%s" % (log_save_folder)) except: # tensorboardX v1.7 train_writer = SummaryWriter(logdir="%s" % (log_save_folder)) time_point = time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime()) textfile = "%s/log_%s.txt" % (log_save_folder, time_point) stdout = Logger(textfile) sys.stdout = stdout print(" ".join(sys.argv)) # initializing parameters for pruning # we can add weights of different layers or we can add gates (multiplies output with 1, useful only for gradient computation) pruning_engine = None if args.pruning: pruning_settings = dict() if not (args.pruning_config is None): pruning_settings_reader = PruningConfigReader() pruning_settings_reader.read_config(args.pruning_config) pruning_settings = pruning_settings_reader.get_parameters() # overwrite parameters from config file with those from command line # needs manual entry here # user_specified = [key for key in vars(default_args).keys() if not (vars(default_args)[key]==vars(args)[key])] # argv_of_interest = ['pruning_threshold', 'pruning-momentum', 'pruning_step', 'prune_per_iteration', # 'fixed_layer', 'start_pruning_after_n_iterations', 'maximum_pruning_iterations', # 'starting_neuron', 'prune_neurons_max', 'pruning_method'] has_attribute = lambda x: any([x in a for a in sys.argv]) if has_attribute('pruning-momentum'): pruning_settings['pruning_momentum'] = vars( args)['pruning_momentum'] if has_attribute('pruning-method'): pruning_settings['method'] = vars(args)['pruning_method'] pruning_parameters_list = prepare_pruning_list( pruning_settings, model, model_name=args.model, pruning_mask_from=args.pruning_mask_from, name=args.name) print("Total pruning layers:", len(pruning_parameters_list)) folder_to_write = "%s" % log_save_folder + "/" log_folder = folder_to_write pruning_engine = pytorch_pruning(pruning_parameters_list, pruning_settings=pruning_settings, log_folder=log_folder) pruning_engine.connect_tensorboard(train_writer) pruning_engine.dataset = args.dataset pruning_engine.model = args.model pruning_engine.pruning_mask_from = args.pruning_mask_from pruning_engine.load_mask() gates_to_params = connect_gates_with_parameters_for_flops( args.model, parameters_for_update_named) pruning_engine.gates_to_params = gates_to_params ###=======================end for pruning # loading model file if (len(args.load_model) > 0) and (not args.dynamic_network): if os.path.isfile(args.load_model): if args.fineturn_model: checkpoint = torch.load(args.load_model) state_dict = checkpoint['state_dict'] model = load_module_state_dict_checkpoint(model, state_dict) else: load_model_pytorch(model, args.load_model, args.model) else: print("=> no checkpoint found at '{}'".format(args.load_model)) exit() if args.tensorboard and 0: if args.dataset == "CIFAR10": dummy_input = torch.rand(1, 3, 32, 32).to(device) elif args.dataset == "Imagenet": dummy_input = torch.rand(1, 3, 224, 224).to(device) train_writer.add_graph(model, dummy_input) for epoch in range(1, args.epochs + 1): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(args, optimizer, epoch, args.zero_lr_for_epochs, train_writer) if not args.run_test and not args.get_inference_time: train(args, model, device, train_loader, optimizer, epoch, criterion, train_writer=train_writer, pruning_engine=pruning_engine) if args.pruning: # skip validation error calculation and model saving if pruning_engine.method == 50: continue # evaluate on validation set prec1 = validate(args, test_loader, model, device, criterion, epoch, train_writer=train_writer) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) model_save_path = "%s/models/checkpoint.weights" % (log_save_folder) paths = "%s/models" % (log_save_folder) model_state_dict = model.state_dict() if args.save_models: save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model_state_dict, 'best_prec1': best_prec1, }, is_best, filename=model_save_path) states = { 'epoch': epoch + 1, 'state_dict': model_state_dict, } torch.save(states, '{}/{}.pth.tar'.format(paths, epoch + 1))
import pandas as pd from PIL import Image from PIL import ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True import os import torch import json import time import argparse import pickle image_transforms = { # Train uses data augmentation 'train': transforms.Compose([ transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)), transforms.RandomRotation(degrees=15), transforms.ColorJitter(), transforms.RandomHorizontalFlip(), transforms.CenterCrop(size=224), # Image net standards transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # Imagenet standards ]), # Validation does not use augmentation 'val': transforms.Compose([ transforms.Resize(size=256), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
def __init__(self, setname, args, augment=False): im_size = args.orig_imsize csv_path = osp.join(SPLIT_PATH, setname + '.csv') cache_path = osp.join( CACHE_PATH, "{}.{}.{}.pt".format(self.__class__.__name__, setname, im_size) ) self.use_im_cache = ( im_size != -1 ) # not using cache if self.use_im_cache: if not osp.exists(cache_path): print('* Cache miss... Preprocessing {}...'.format(setname)) resize_ = identity if im_size < 0 else transforms.Resize(im_size) data, label = self.parse_csv(csv_path, setname) self.data = [ resize_(Image.open(path).convert('RGB')) for path in data ] self.label = label print('* Dump cache from {}'.format(cache_path)) torch.save({'data': self.data, 'label': self.label }, cache_path) else: print('* Load cache from {}'.format(cache_path)) cache = torch.load(cache_path) self.data = cache['data'] self.label = cache['label'] else: self.data, self.label = self.parse_csv(csv_path, setname) self.num_class = len(set(self.label)) image_size = 84 if augment and setname == 'train': transforms_list = [ transforms.RandomResizedCrop(image_size), transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ] else: transforms_list = [ transforms.Resize(92), transforms.CenterCrop(image_size), transforms.ToTensor(), ] # Transformation if args.backbone_class == 'ConvNet': self.transform = transforms.Compose( transforms_list + [ transforms.Normalize(np.array([0.485, 0.456, 0.406]), np.array([0.229, 0.224, 0.225])) ]) elif args.backbone_class == 'Res12': self.transform = transforms.Compose( transforms_list + [ transforms.Normalize(np.array([x / 255.0 for x in [120.39586422, 115.59361427, 104.54012653]]), np.array([x / 255.0 for x in [70.68188272, 68.27635443, 72.54505529]])) ]) elif args.backbone_class == 'Res18': self.transform = transforms.Compose( transforms_list + [ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) elif args.backbone_class == 'WRN': self.transform = transforms.Compose( transforms_list + [ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) else: raise ValueError('Non-supported Network Types. Please Revise Data Pre-Processing Scripts.')
def load_data(train=True, data_dir='dataset/imagenet', batch_size=128, subset_len=None, sample_method='random', distributed=False, model_name='resnet18', **kwargs): #prepare data # random.seed(12345) traindir = data_dir + '/train' valdir = data_dir + '/val' train_sampler = None normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if model_name == 'inception_v3': size = 299 resize = 299 else: size = 224 resize = 256 if train: dataset = torchvision.datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if subset_len: assert subset_len <= len(dataset) if sample_method == 'random': dataset = torch.utils.data.Subset( dataset, random.sample(range(0, len(dataset)), subset_len)) else: dataset = torch.utils.data.Subset(dataset, list(range(subset_len))) if distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( dataset) data_loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=(train_sampler is None), sampler=train_sampler, **kwargs) else: dataset = torchvision.datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(resize), transforms.CenterCrop(size), transforms.ToTensor(), normalize, ])) if subset_len: assert subset_len <= len(dataset) if sample_method == 'random': dataset = torch.utils.data.Subset( dataset, random.sample(range(0, len(dataset)), subset_len)) else: dataset = torch.utils.data.Subset(dataset, list(range(subset_len))) data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, **kwargs) return data_loader, train_sampler
def train(args): # Create model directory if not os.path.exists(args.model_path): os.makedirs(args.model_path) # Save the arguments. with open(os.path.join(args.model_path, 'args.json'), 'w') as args_file: json.dump(args.__dict__, args_file) # Config logging. log_format = '%(levelname)-8s %(message)s' logfile = os.path.join(args.model_path, 'train.log') logging.basicConfig(filename=logfile, level=logging.INFO, format=log_format) logging.getLogger().addHandler(logging.StreamHandler()) logging.info(json.dumps(args.__dict__)) # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.ToPILImage(), transforms.RandomResizedCrop(args.crop_size, scale=(1.00, 1.2), ratio=(0.75, 1.3333333333333333)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])##0.5 # Load vocabulary wrapper. vocab = load_vocab(args.vocab_path) # Build data loader logging.info("Building data loader...") train_sampler = None val_sampler = None data_loader = get_loader(args.dataset, transform, args.batch_size, shuffle=False, num_workers=args.num_workers, max_examples=args.max_examples, sampler=train_sampler) val_data_loader = get_loader(args.val_dataset, transform, args.batch_size, shuffle=False, num_workers=args.num_workers, max_examples=args.max_examples, sampler=val_sampler) logging.info("Done") vqg = create_model(args, vocab) print(vqg) """ graph = hl.build_graph(vqg.modules()) graph = graph.build_dot() graph.render("weights/tf1/", view=True, format='png') g = make_dot(vqg.modules(),vqg.state_dict()) g.view()""" if args.load_model is not None: vqg.load_state_dict(torch.load(args.load_model)) logging.info("Done") # Loss criterion. pad = vocab(vocab.SYM_PAD) # Set loss weight for 'pad' symbol to 0 criterion = nn.CrossEntropyLoss(ignore_index=pad) l2_criterion = nn.MSELoss() # Setup GPUs. if torch.cuda.is_available(): logging.info("Using available GPU...") vqg.cuda() criterion.cuda() l2_criterion.cuda() # Parameters to train. gen_params = vqg.generator_parameters() info_params = vqg.info_parameters() learning_rate = args.learning_rate info_learning_rate = args.learning_rate gen_optimizer = torch.optim.Adam(gen_params, lr=learning_rate) info_optimizer = torch.optim.Adam(info_params, lr=info_learning_rate) scheduler = ReduceLROnPlateau(optimizer=gen_optimizer, mode='min', factor=0.1, patience=args.patience, verbose=True, min_lr=1e-7) info_scheduler = ReduceLROnPlateau(optimizer=info_optimizer, mode='min', factor=0.1, patience=args.patience, verbose=True, min_lr=1e-7) # Train the model. total_steps = len(data_loader) start_time = time.time() n_steps = 0 for epoch in range(args.num_epochs): for i, (images, questions, qindices) in enumerate(data_loader): n_steps += 1 # Set mini-batch dataset. if torch.cuda.is_available(): images = images.cuda() questions = questions.cuda() qindices = qindices.cuda() # Eval now. if (args.eval_every_n_steps is not None and n_steps >= args.eval_every_n_steps and n_steps % args.eval_every_n_steps == 0): run_eval(vqg, val_data_loader, criterion, l2_criterion, args, epoch, scheduler,info_scheduler) compare_outputs(images, questions, vqg, vocab, logging, args) # Forward. vqg.train() gen_optimizer.zero_grad() info_optimizer.zero_grad() image_features = vqg.encode_images(images) # Question generation. mus, logvars = vqg.encode_into_z(image_features) zs = vqg.reparameterize(mus, logvars) (outputs, _, _) = vqg.decode_questions( image_features, zs, questions=questions, teacher_forcing_ratio=1.0) # Reorder the questions based on length. questions = torch.index_select(questions, 0, qindices) # Ignoring the start token. questions = questions[:, 1:] qlengths = process_lengths(questions) # Convert the output from MAX_LEN list of (BATCH x VOCAB) -> # (BATCH x MAX_LEN x VOCAB). outputs = [o.unsqueeze(1) for o in outputs] outputs = torch.cat(outputs, dim=1) outputs = torch.index_select(outputs, 0, qindices) # Calculate the generation loss. targets = pack_padded_sequence(questions, qlengths, batch_first=True)[0] outputs = pack_padded_sequence(outputs, qlengths, batch_first=True)[0] gen_loss = criterion(outputs, targets) total_loss = 0.0 total_loss += args.lambda_gen * gen_loss gen_loss = gen_loss.item() # Variational loss. kl_loss = gaussian_KL_loss(mus, logvars) total_loss += args.lambda_z * kl_loss kl_loss = kl_loss.item() total_loss.backward() gen_optimizer.step() # Print log info if i % args.log_step == 0: delta_time = time.time() - start_time start_time = time.time() logging.info('Time: %.4f, Epoch [%d/%d], Step [%d/%d], ' 'LR: %f, gen: %.4f, KL: %.4f ' % (delta_time, epoch, args.num_epochs, i, total_steps, gen_optimizer.param_groups[0]['lr'], gen_loss, kl_loss )) # Save the models if args.save_step is not None and (i+1) % args.save_step == 0: torch.save(vqg.state_dict(), os.path.join(args.model_path, 'vqg-tf-%d-%d.pkl' % (epoch + 1, i + 1))) torch.save(vqg.state_dict(), os.path.join(args.model_path, 'vqg-tf-%d.pkl' % (epoch+1))) # Evaluation and learning rate updates. run_eval(vqg, val_data_loader, criterion, l2_criterion, args, epoch, scheduler,info_scheduler)
def main(): global args, best_prec1 args = parser.parse_args() if not os.path.isdir(args.checkpoint): mkdir_p(args.checkpoint) print('dataset:', args.dataset) if args.dataset == 'EPIC': with open(args.foptions, 'rb') as handle: options = pickle.load(handle) if options['root'][0] == '.': # update relative path options['root'] = 'basemodel' + options['root'][1:] options['num_crops'] = 1 else: options = None model = models.Net(options=options).cuda() # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() extractor_params = list(map(id, model.extractor.parameters())) classifier_params = filter(lambda p: id(p) not in extractor_params, model.parameters()) optimizer = torch.optim.SGD([{ 'params': model.extractor.parameters() }, { 'params': classifier_params, 'lr': args.lr * 10 }], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1) # optionally resume from a checkpoint # title = 'CUB' title = args.dataset if args.resume: print('==> Resuming from checkpoint..') assert os.path.isfile( args.resume), 'Error: no checkpoint directory found!' checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title) logger.set_names([ 'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.' ]) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) # Data loading code normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) if args.dataset == 'CUB': print('Using dataset CUB') train_dataset = loader.ImageLoader( args.data, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]), train=True) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(loader.ImageLoader( args.data, 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) elif args.dataset == 'EPIC': print('Using dataset EPIC') train_dataset, val_dataset, train_loader, val_loader = get_datasets_and_dataloaders( options, device=options['device']) if args.evaluate: validate(val_loader, model, criterion) return print('# train_loader:', len(train_loader)) print('# val_loader:', len(val_loader)) print('Training epoch from {} to {}'.format(args.start_epoch, args.epochs)) for epoch in range(args.start_epoch, args.epochs): scheduler.step() lr = optimizer.param_groups[1]['lr'] print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, lr)) # train for one epoch train_loss, train_acc = train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set test_loss, test_acc = validate(val_loader, model, criterion) # append logger file logger.append([lr, train_loss, test_loss, train_acc, test_acc]) # remember best prec@1 and save checkpoint is_best = test_acc > best_prec1 best_prec1 = max(test_acc, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, is_best, checkpoint=args.checkpoint) logger.close() logger.plot() savefig(os.path.join(args.checkpoint, 'log.eps')) print('Best acc:') print(best_prec1)
def main_worker(gpu, ngpus_per_node, args): global best_acc1 # global aim_sess = aim.Session() # aim_sess.set_params({ # 'num_epochs': args.epochs, # 'num_classes': 1000, # 'batch_size': args.batch_size, # }, name='hparams') 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) repvgg_build_func = get_RepVGG_func_by_name(args.arch) model = repvgg_build_func(deploy=False) 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 model = torch.nn.DataParallel(model).cuda() # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda(args.gpu) optimizer = sgd_optimizer(model, args.lr, args.momentum, args.weight_decay) lr_scheduler = CosineAnnealingLR( optimizer=optimizer, T_max=args.epochs * IMAGENET_TRAINSET_SIZE // args.batch_size // ngpus_per_node) # 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']) lr_scheduler.load_state_dict(checkpoint['scheduler']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code 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_dataset = datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])) val_loader = torch.utils.data.DataLoader(val_dataset, 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, lr_scheduler) # 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(), 'scheduler': lr_scheduler.state_dict(), }, is_best)
def build_imagenet(model_state_dict, optimizer_state_dict, **kwargs): ratio = kwargs.pop('ratio') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.2), transforms.ToTensor(), normalize, ]) valid_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]) if args.zip_file: logging.info('Loading data from zip file') traindir = os.path.join(args.data, 'train.zip') if args.lazy_load: train_data = utils.ZipDataset(traindir, train_transform) else: logging.info('Loading data into memory') train_data = utils.InMemoryZipDataset(traindir, train_transform, num_workers=32) else: logging.info('Loading data from directory') traindir = os.path.join(args.data, 'train') if args.lazy_load: train_data = dset.ImageFolder(traindir, train_transform) else: logging.info('Loading data into memory') train_data = utils.InMemoryDataset(traindir, train_transform, num_workers=32) num_train = len(train_data) indices = list(range(num_train)) np.random.shuffle(indices) split = int(np.floor(ratio * num_train)) train_indices = sorted(indices[:split]) valid_indices = sorted(indices[split:]) train_queue = torch.utils.data.DataLoader( train_data, batch_size=args.child_batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices), pin_memory=True, num_workers=16) valid_queue = torch.utils.data.DataLoader( train_data, batch_size=args.child_eval_batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_indices), pin_memory=True, num_workers=16) model = NASWSNetworkImageNet(args, 1000, args.child_layers, args.child_nodes, args.child_channels, args.child_keep_prob, args.child_drop_path_keep_prob, args.child_use_aux_head, args.steps) model = model.cuda() train_criterion = CrossEntropyLabelSmooth(1000, args.child_label_smooth).cuda() eval_criterion = nn.CrossEntropyLoss().cuda() logging.info("param size = %fMB", utils.count_parameters_in_MB(model)) optimizer = torch.optim.SGD( model.parameters(), args.child_lr, momentum=0.9, weight_decay=args.child_l2_reg, ) if model_state_dict is not None: model.load_state_dict(model_state_dict) if optimizer_state_dict is not None: optimizer.load_state_dict(optimizer_state_dict) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.child_decay_period, gamma=args.child_gamma) return train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler
def main(): parser = argparse.ArgumentParser() parser.add_argument('--dataroot', default='/mnt/Data1', help='path to images') parser.add_argument('--workers', default=4, type=int, help='number of data loading workers (default: 4)') parser.add_argument('--batch-size', default=2, type=int, help='mini-batch size') parser.add_argument('--outroot', default='./results', help='path to save the results') parser.add_argument('--exp-name', default='test', help='name of expirement') parser.add_argument('--load', default='', help='name of pth to load weights from') parser.add_argument('--freeze-cc-net', dest='freeze_cc_net', action='store_true', help='dont train the color corrector net') parser.add_argument('--freeze-warp-net', dest='freeze_warp_net', action='store_true', help='dont train the warp net') parser.add_argument('--test', dest='test', action='store_true', help='only test the network') parser.add_argument( '--synth-data', dest='synth_data', action='store_true', help='use synthetic data instead of tank data for training') parser.add_argument('--epochs', default=3, type=int, help='number of epochs to train for') parser.add_argument('--no-warp-net', dest='warp_net', action='store_false', help='do not include warp net in the model') parser.add_argument('--warp-net-downsample', default=3, type=int, help='number of downsampling layers in warp net') parser.add_argument('--no-color-net', dest='color_net', action='store_false', help='do not include color net in the model') parser.add_argument('--color-net-downsample', default=3, type=int, help='number of downsampling layers in color net') parser.add_argument( '--no-color-net-skip', dest='color_net_skip', action='store_false', help='dont use u-net skip connections in the color net') parser.add_argument( '--dim', default=32, type=int, help='initial feature dimension (doubled at each downsampling layer)') parser.add_argument('--n-res', default=8, type=int, help='number of residual blocks') parser.add_argument('--norm', default='gn', type=str, help='type of normalization layer') parser.add_argument( '--denormalize', dest='denormalize', action='store_true', help='denormalize output image by input image mean/var') parser.add_argument( '--weight-X-L1', default=1., type=float, help='weight of L1 reconstruction loss after color corrector net') parser.add_argument('--weight-Y-L1', default=1., type=float, help='weight of L1 reconstruction loss after warp net') parser.add_argument('--weight-Y-VGG', default=1., type=float, help='weight of perceptual loss after warp net') parser.add_argument( '--weight-Z-L1', default=1., type=float, help='weight of L1 reconstruction loss after color net') parser.add_argument('--weight-Z-VGG', default=.5, type=float, help='weight of perceptual loss after color net') parser.add_argument('--weight-Z-Adv', default=0.2, type=float, help='weight of adversarial loss after color net') args = parser.parse_args() # set random seed for consistent fixed batch torch.manual_seed(8) # set weights of losses of intermediate outputs to zero if not necessary if not args.warp_net: args.weight_Y_L1 = 0 args.weight_Y_VGG = 0 if not args.color_net: args.weight_Z_L1 = 0 args.weight_Z_VGG = 0 args.weight_Z_Adv = 0 # datasets train_dir_1 = os.path.join(args.dataroot, 'Water', 'train') train_dir_2 = os.path.join(args.dataroot, 'ImageNet', 'train') val_dir_1 = os.path.join(args.dataroot, 'Water', 'test') val_dir_2 = os.path.join(args.dataroot, 'ImageNet', 'test') test_dir = os.path.join(args.dataroot, 'Water_Real') if args.synth_data: train_data = ImageFolder(train_dir_2, transform=transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), synthdata.SynthData(224, n=args.batch_size), ])) else: train_data = PairedImageFolder( train_dir_1, train_dir_2, transform=transforms.Compose([ pairedtransforms.RandomResizedCrop(224), pairedtransforms.RandomHorizontalFlip(), pairedtransforms.ToTensor(), pairedtransforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ])) val_data = PairedImageFolder(val_dir_1, val_dir_2, transform=transforms.Compose([ pairedtransforms.Resize(256), pairedtransforms.CenterCrop(256), pairedtransforms.ToTensor(), pairedtransforms.Normalize( mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ])) test_data = ImageFolder(test_dir, transform=transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(256), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]), return_path=True) train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, num_workers=args.workers, pin_memory=True, shuffle=True) val_loader = torch.utils.data.DataLoader(val_data, batch_size=args.batch_size, num_workers=args.workers, pin_memory=True, shuffle=True) test_loader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, num_workers=args.workers, pin_memory=True, shuffle=False) # fixed test batch for visualization during training fixed_batch = iter(val_loader).next()[0] # model model = networks.Model(args) model.cuda() # load weights from checkpoint if args.test and not args.load: args.load = args.exp_name if args.load: model.load_state_dict(torch.load( os.path.join(args.outroot, '%s_net.pth' % args.load)), strict=args.test) # create outroot if necessary if not os.path.exists(args.outroot): os.makedirs(args.outroot) # if args.test only run test script if args.test: test(test_loader, model, args) return # main training loop for epoch in range(args.epochs): train(train_loader, model, fixed_batch, epoch, args) torch.save(model.state_dict(), os.path.join(args.outroot, '%s_net.pth' % args.exp_name)) test(test_loader, model, args)
def transforming(path): # Transform image data into tensors to be used in PyTorch # Getting the specific paths to each of the training, validation and testing folder data_dir = path train_dir = data_dir + '/train' valid_dir = data_dir + '/valid' test_dir = data_dir + '/test' # Preparing the transformations needed to convert images into tensors data_transforms = { 'training': transforms.Compose([ transforms.RandomRotation(30), transforms.RandomHorizontalFlip(30), transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'validation': transforms.Compose([ transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'testing': transforms.Compose([ transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } # Loading the data from the directories image_datasets = { 'training': datasets.ImageFolder(train_dir, transform=data_transforms['training']), 'validation': datasets.ImageFolder(valid_dir, transform=data_transforms['validation']), 'testing': datasets.ImageFolder(test_dir, transform=data_transforms['testing']) } # Batch iterators to feed the network dataloaders = { 'training': torch.utils.data.DataLoader(image_datasets['training'], batch_size=100, shuffle=True), 'validation': torch.utils.data.DataLoader(image_datasets['validation'], batch_size=100, shuffle=True), 'testing': torch.utils.data.DataLoader(image_datasets['testing'], batch_size=100, shuffle=True) } return dataloaders, image_datasets
sample = np.array(sample) # blur the image with a 50% chance prob = np.random.random_sample() if prob < 0.5: sigma = (self.max - self.min) * np.random.random_sample() + self.min sample = cv2.GaussianBlur(sample, (self.kernel_size, self.kernel_size), sigma) return sample train_transform = transforms.Compose([ transforms.RandomResizedCrop(32), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8), transforms.RandomGrayscale(p=0.2), GaussianBlur(kernel_size=int(0.1 * 32)), transforms.ToTensor(), transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]) ]) test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]) ])
def __init__(self, normalize=None, size_crops: List[int] = [96, 36], nmb_crops: List[int] = [2, 4], min_scale_crops: List[float] = [0.33, 0.10], max_scale_crops: List[float] = [1, 0.33], gaussian_blur: bool = True, jitter_strength: float = 1.): self.jitter_strength = jitter_strength self.gaussian_blur = gaussian_blur assert len(size_crops) == len(nmb_crops) assert len(min_scale_crops) == len(nmb_crops) assert len(max_scale_crops) == len(nmb_crops) self.size_crops = size_crops self.nmb_crops = nmb_crops self.min_scale_crops = min_scale_crops self.max_scale_crops = max_scale_crops self.color_jitter = transforms.ColorJitter(0.8 * self.jitter_strength, 0.8 * self.jitter_strength, 0.8 * self.jitter_strength, 0.2 * self.jitter_strength) transform = [] color_transform = [ transforms.RandomApply([self.color_jitter], p=0.8), transforms.RandomGrayscale(p=0.2) ] if self.gaussian_blur: kernel_size = int(0.1 * self.size_crops[0]) if kernel_size % 2 == 0: kernel_size += 1 color_transform.append(GaussianBlur(kernel_size=kernel_size, p=0.5)) self.color_transform = transforms.Compose(color_transform) if normalize is None: self.final_transform = transforms.ToTensor() else: self.final_transform = transforms.Compose( [transforms.ToTensor(), normalize]) for i in range(len(self.size_crops)): random_resized_crop = transforms.RandomResizedCrop( self.size_crops[i], scale=(self.min_scale_crops[i], self.max_scale_crops[i]), ) transform.extend([ transforms.Compose([ random_resized_crop, transforms.RandomHorizontalFlip(p=0.5), self.color_transform, self.final_transform ]) ] * self.nmb_crops[i]) self.transform = transform # add online train transform of the size of global view online_train_transform = transforms.Compose([ transforms.RandomResizedCrop(self.size_crops[0]), transforms.RandomHorizontalFlip(), self.final_transform ]) self.transform.append(online_train_transform)
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 = models.__dict__[args.arch]() n_ftrs = model.fc.in_features model.fc = nn.Linear(n_ftrs, 3) 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) 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)) checkpoint = torch.load(args.resume) 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 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 = Synthetic('../annotation.json', subset='training', transform=transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_dataset = Synthetic('../annotation.json', subset='validation', transform=transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), 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(val_dataset, 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)
def load_data( prop_noPV_training: float, min_rescale_images: float, batch_size: int, dir_data_training: str = "", dir_data_validation: str = "", dir_data_test: str = "", ): """ Create the DataLoader objects that will generate the training, validation and test sets. Parameters ---------- prop_noPV_training : float Proportion of noPV images to add for the training of the model. min_rescale_images : float Minimum proportion of the image to keep for the RandomResizedCrop transform. batch_size : int Number of samples per batch in the DataLoaders. dir_data_training : str, optional Directory where the folders "images/", "labels/" and "noPV/" are for the training set. If empty, the data is not generated. The default is "". dir_data_validation : str, optional Directory where the folders "images/", "labels/" and "noPV/" are for the validation set. If empty, the data is not generated. The default is "". dir_data_test : str, optional Directory where the folders "images/", "labels/" and "noPV/" are for the test set. If empty, the data is not generated. The default is "". Returns ------- train_dl : torch.utils.data.DataLoader Training DataLoader, if data directory is provided, otherwise None. validation_dl : torch.utils.data.DataLoader Validation DataLoader, if data directory is provided, otherwise None. test_dl : torch.utils.data.DataLoader Test DataLoader, if data directory is provided, otherwise None. """ roof_train_dataset = None if dir_data_training: # Transforms to augment the data (for training set) transform_aug = transforms.Compose([ transforms.ToPILImage(), transforms.RandomResizedCrop(250, scale=(min_rescale_images, 1.0), ratio=(1.0, 1.0)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) # Instantiate the training dataset roof_train_dataset = AvailableRooftopDataset( dir_PV=os.path.join(dir_data_training, "PV"), dir_noPV=os.path.join(dir_data_training, "noPV"), dir_labels=os.path.join(dir_data_training, "labels"), transform=transform_aug, prop_noPV=prop_noPV_training, ) # No transform applied to validation and train images (the model should not need # any preprocessing) transform_id = transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), ]) # Instantiate the validation and test datasets roof_validation_dataset, roof_test_dataset = ( AvailableRooftopDataset( dir_PV=os.path.join(dir_data, "PV"), dir_noPV=os.path.join(dir_data, "noPV"), dir_labels=os.path.join(dir_data, "labels"), transform=transform_id, prop_noPV=-1, # All of them ) if dir_data else None for dir_data in (dir_data_validation, dir_data_test)) # Instantiate the DataLoaders roof_train_dl, roof_validation_dl, roof_test_dl = (DataLoader( roof_dataset, batch_size=batch_size, shuffle=True, num_workers=0) if roof_dataset else None for roof_dataset in ( roof_train_dataset, roof_validation_dataset, roof_test_dataset, )) return roof_train_dl, roof_validation_dl, roof_test_dl
def get_dataloader(args, add_erasing=False, aug_plus=False): if 'cifar' in args.dataset or 'kitchen' in args.dataset: if aug_plus: # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709 transform_train_list = [ transforms.RandomResizedCrop(size=32, scale=(0.2, 1.)), transforms.RandomApply( [ transforms.ColorJitter(0.4, 0.4, 0.4, 0.1) # not strengthened ], p=0.8), transforms.RandomGrayscale(p=0.2), transforms.RandomApply([GaussianBlur([.1, 2.])], p=0.5), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ] else: transform_train_list = [ transforms.RandomResizedCrop(size=32, scale=(0.2, 1.)), transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), transforms.RandomGrayscale(p=0.2), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ] if add_erasing: transform_train_list.append(transforms.RandomErasing(p=1.0)) transform_train = transforms.Compose(transform_train_list) if 'kitchen' in args.dataset: transform_test = transforms.Compose([ transforms.Resize((32, 32), interpolation=2), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) else: transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) elif 'stl' in args.dataset: transform_train = transforms.Compose([ transforms.RandomResizedCrop(size=96, scale=(0.2, 1.)), transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), transforms.RandomGrayscale(p=0.2), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) if args.dataset == 'cifar10': trainset = datasets.CIFAR10Instance(root='./data/CIFAR-10', train=True, download=True, transform=transform_train, two_imgs=args.two_imgs, three_imgs=args.three_imgs) train_sampler = torch.utils.data.distributed.DistributedSampler( trainset) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, sampler=train_sampler) testset = datasets.CIFAR10Instance(root='./data/CIFAR-10', train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2, pin_memory=False) args.pool_len = 4 ndata = trainset.__len__() elif args.dataset == 'cifar100': trainset = datasets.CIFAR100Instance(root='./data/CIFAR-100', train=True, download=True, transform=transform_train, two_imgs=args.two_imgs, three_imgs=args.three_imgs) train_sampler = torch.utils.data.distributed.DistributedSampler( trainset) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, sampler=train_sampler) testset = datasets.CIFAR100Instance(root='./data/CIFAR-100', train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2, pin_memory=False) args.pool_len = 4 ndata = trainset.__len__() elif args.dataset == 'stl10': trainset = datasets.STL10(root='./data/STL10', split='train', download=True, transform=transform_train, two_imgs=args.two_imgs, three_imgs=args.three_imgs) train_sampler = torch.utils.data.distributed.DistributedSampler( trainset) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, sampler=train_sampler) testset = datasets.STL10(root='./data/STL10', split='test', download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2, pin_memory=False) args.pool_len = 7 ndata = trainset.__len__() elif args.dataset == 'stl10-full': trainset = datasets.STL10(root='./data/STL10', split='train+unlabeled', download=True, transform=transform_train, two_imgs=args.two_imgs, three_imgs=args.three_imgs) train_sampler = torch.utils.data.distributed.DistributedSampler( trainset) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, sampler=train_sampler) labeledTrainset = datasets.STL10(root='./data/STL10', split='train', download=True, transform=transform_train, two_imgs=args.two_imgs) labeledTrain_sampler = torch.utils.data.distributed.DistributedSampler( labeledTrainset) labeledTrainloader = torch.utils.data.DataLoader( labeledTrainset, batch_size=args.batch_size, shuffle=False, num_workers=2, pin_memory=False, sampler=labeledTrain_sampler) testset = datasets.STL10(root='./data/STL10', split='test', download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2, pin_memory=False) args.pool_len = 7 ndata = labeledTrainset.__len__() elif args.dataset == 'kitchen': trainset = datasets.CIFARImageFolder(root='./data/Kitchen-HC/train', train=True, transform=transform_train, two_imgs=args.two_imgs, three_imgs=args.three_imgs) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=False) testset = datasets.CIFARImageFolder(root='./data/Kitchen-HC/test', train=False, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2, pin_memory=False) args.pool_len = 4 ndata = trainset.__len__() return trainloader, testloader, ndata
def main(): global best_acc1 best_acc1 = 0 args = parse_option() if args.gpu is not None: print("Use GPU: {} for training".format(args.gpu)) # set the data loader train_folder = os.path.join(args.data_folder, 'train') val_folder = os.path.join(args.data_folder, 'val') image_size = 224 crop_padding = 32 mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] normalize = transforms.Normalize(mean=mean, std=std) if args.aug == 'NULL': train_transform = transforms.Compose([ transforms.RandomResizedCrop(image_size, scale=(args.crop, 1.)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) elif args.aug == 'CJ': train_transform = transforms.Compose([ transforms.RandomResizedCrop(image_size, scale=(args.crop, 1.)), transforms.RandomGrayscale(p=0.2), transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) else: raise NotImplemented('augmentation not supported: {}'.format(args.aug)) train_dataset = datasets.ImageFolder(train_folder, train_transform) val_dataset = datasets.ImageFolder( val_folder, transforms.Compose([ transforms.Resize(image_size + crop_padding), transforms.CenterCrop(image_size), transforms.ToTensor(), normalize, ])) print(len(train_dataset)) train_sampler = None train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.num_workers, pin_memory=True, sampler=train_sampler) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) # create model and optimizer if args.model == 'resnet50': model = InsResNet50() classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 1) elif args.model == 'resnet50x2': model = InsResNet50(width=2) classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 2) elif args.model == 'resnet50x4': model = InsResNet50(width=4) classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 4) else: raise NotImplementedError('model not supported {}'.format(args.model)) print('==> loading pre-trained model') ckpt = torch.load(args.model_path) model.load_state_dict(ckpt['model']) print("==> loaded checkpoint '{}' (epoch {})".format( args.model_path, ckpt['epoch'])) print('==> done') model = model.cuda() classifier = classifier.cuda() criterion = torch.nn.CrossEntropyLoss().cuda(args.gpu) if not args.adam: optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(args.beta1, args.beta2), weight_decay=args.weight_decay, eps=1e-8) model.eval() cudnn.benchmark = True # set mixed precision training # if args.amp: # model = amp.initialize(model, opt_level=args.opt_level) # classifier, optimizer = amp.initialize(classifier, optimizer, opt_level=args.opt_level) # optionally resume from a checkpoint args.start_epoch = 1 if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume, map_location='cpu') # checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] + 1 classifier.load_state_dict(checkpoint['classifier']) optimizer.load_state_dict(checkpoint['optimizer']) best_acc1 = checkpoint['best_acc1'] best_acc1 = best_acc1.cuda() print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) if 'opt' in checkpoint.keys(): # resume optimization hyper-parameters print('=> resume hyper parameters') if 'bn' in vars(checkpoint['opt']): print('using bn: ', checkpoint['opt'].bn) if 'adam' in vars(checkpoint['opt']): print('using adam: ', checkpoint['opt'].adam) if 'cosine' in vars(checkpoint['opt']): print('using cosine: ', checkpoint['opt'].cosine) args.learning_rate = checkpoint['opt'].learning_rate # args.lr_decay_epochs = checkpoint['opt'].lr_decay_epochs args.lr_decay_rate = checkpoint['opt'].lr_decay_rate args.momentum = checkpoint['opt'].momentum args.weight_decay = checkpoint['opt'].weight_decay args.beta1 = checkpoint['opt'].beta1 args.beta2 = checkpoint['opt'].beta2 del checkpoint torch.cuda.empty_cache() else: print("=> no checkpoint found at '{}'".format(args.resume)) # set cosine annealing scheduler if args.cosine: # last_epoch = args.start_epoch - 2 # eta_min = args.learning_rate * (args.lr_decay_rate ** 3) * 0.1 # scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs, eta_min, last_epoch) eta_min = args.learning_rate * (args.lr_decay_rate**3) * 0.1 scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, args.epochs, eta_min, -1) # dummy loop to catch up with current epoch for i in range(1, args.start_epoch): scheduler.step() # tensorboard logger = tb_logger.Logger(logdir=args.tb_folder, flush_secs=2) # routine for epoch in range(args.start_epoch, args.epochs + 1): if args.cosine: scheduler.step() else: adjust_learning_rate(epoch, args, optimizer) print("==> training...") time1 = time.time() train_acc, train_acc5, train_loss = train(epoch, train_loader, model, classifier, criterion, optimizer, args) time2 = time.time() print('train epoch {}, total time {:.2f}'.format(epoch, time2 - time1)) logger.log_value('train_acc', train_acc, epoch) logger.log_value('train_acc5', train_acc5, epoch) logger.log_value('train_loss', train_loss, epoch) logger.log_value('learning_rate', optimizer.param_groups[0]['lr'], epoch) print("==> testing...") test_acc, test_acc5, test_loss = validate(val_loader, model, classifier, criterion, args) logger.log_value('test_acc', test_acc, epoch) logger.log_value('test_acc5', test_acc5, epoch) logger.log_value('test_loss', test_loss, epoch) # save the best model if test_acc > best_acc1: best_acc1 = test_acc state = { 'opt': args, 'epoch': epoch, 'classifier': classifier.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), } save_name = '{}_layer{}.pth'.format(args.model, args.layer) save_name = os.path.join(args.save_folder, save_name) print('saving best model!') torch.save(state, save_name) # save model if epoch % args.save_freq == 0: print('==> Saving...') state = { 'opt': args, 'epoch': epoch, 'classifier': classifier.state_dict(), 'best_acc1': test_acc, 'optimizer': optimizer.state_dict(), } save_name = 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch) save_name = os.path.join(args.save_folder, save_name) print('saving regular model!') torch.save(state, save_name) # tensorboard logger pass
def main_worker(gpu, ngpus_per_node, args): global best_acc1 args.gpu = gpu lr_decay_epoch = [int(i) for i in args.lr_decay_epoch.split(',')] 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 = models.__dict__[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() criterion = nn.CrossEntropyLoss().cuda(args.gpu) logname = args.prefix if args.optimizer == 'sgd': logname += "SGD_" print("sgd") optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'nero': logname += "Nero_" print("Nero") optimizer = Nero(model.parameters(), lr=args.lr, constraints=True) cos_sch = False scheduler = None T_max = math.ceil(1281167.0 / float(args.batch_size)) * (args.epochs) if args.sch == 'cos' or args.sch == 'cosine': print("cosine scheduler") cos_sch = True scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=T_max, eta_min=0.0) logname += args.arch + "_sch_" +str(args.sch)+ "_lr" +str(args.lr) + \ '_epoch' + str(args.epochs) + \ "_opt_" + args.optimizer + \ "_b" + str(args.batch_size) + \ '_momentum' + str(args.momentum) + "_beta" + str(args.beta) + \ '_wd' + str(args.weight_decay) writer = SummaryWriter(args.logdir + '/' + logname) # 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'], strict=False) optimizer.load_state_dict(checkpoint['optimizer']) for group in optimizer.param_groups: group["lr"] = args.lr if args.sch == 'cos': for i in range(checkpoint['epoch'] * math.ceil(1281167.0 / float(args.batch_size))): scheduler.step() print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code 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 lr = args.lr for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) current_lr = optimizer.param_groups[0]['lr'] print("current learning rate: {}".format(current_lr)) writer.add_scalar('lr', current_lr, epoch) # train for one epoch top1_train, top5_train, losses_train, batch_time_train, scheduler = train( train_loader, model, criterion, optimizer, epoch, args, writer, scheduler=scheduler) if not cos_sch: lr = adjust_learning_rate(optimizer, epoch, lr, lr_decay_epoch, args.lr_decay) writer.add_scalar('train/batch_time_mean', batch_time_train, epoch) writer.add_scalar('train/loss_mean', losses_train, epoch) writer.add_scalar('train/top1_mean', top1_train, epoch) writer.add_scalar('train/top5_mean', top5_train, epoch) # evaluate on validation set top1_val, top5_val, losses_val, batch_time_val = validate( val_loader, model, criterion, args, epoch, writer) writer.add_scalar('val/batch_time_mean', batch_time_val, epoch) writer.add_scalar('val/loss_mean', losses_val, epoch) writer.add_scalar('val/top1_mean', top1_val, epoch) writer.add_scalar('val/top5_mean', top5_val, epoch) # remember best acc@1 and save checkpoint is_best = top1_val > best_acc1 best_acc1 = max(top1_val, best_acc1) if not args.multiprocessing_distributed or ( args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): if ((epoch + 1) % 5 == 0): save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), }, is_best, args.logdir + '/' + logname + '/epoch' + str(epoch + 1) + '_checkpoint.pth.tar') writer.close()
import torch import torchvision.transforms as transforms from PIL import Image from torch.autograd import Variable # Defines where the model is located. model_path = 'models/classify_resnet_152_80.pth' # Defines where the test data is located. test_folder = './data/test' # Defines how often to print progress. print_batch_size = 50 # Defines how to pre-process the image data. transform = transforms.Compose([ transforms.RandomResizedCrop(200), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # Loads the model and sets to eval mode. model = torch.load(model_path) model.eval() # Iterates over all images in test folder. output = [] for index, row in pd.read_csv('./data/test.csv').iterrows(): if index % print_batch_size == 0: print('Progress: #{:5d} | time: {}'.format(index, time.ctime()))
def get_dataloader(args): # Define data files path. root_img_folder = "/ais/gobi4/fashion/edge_detection/data_aug" root_label_folder = "/ais/gobi4/fashion/edge_detection/data_aug" train_anno_txt = "/ais/gobi4/fashion/edge_detection/data_aug/list_train_aug.txt" val_anno_txt = "/ais/gobi4/fashion/edge_detection/data_aug/list_test.txt" train_hdf5_file = "/ais/gobi6/jiaman/github/CASENet/utils/train_aug_label_binary_np.h5" val_hdf5_file = "/ais/gobi6/jiaman/github/CASENet/utils/test_label_binary_np.h5" input_size = 472 normalize = transforms.Normalize(mean=[104.008, 116.669, 122.675], std=[1, 1, 1]) train_augmentation = transforms.Compose([ transforms.RandomResizedCrop(input_size, scale=(0.75, 1.0), ratio=(0.75, 1.0)), transforms.RandomHorizontalFlip() ]) train_label_augmentation = transforms.Compose([transforms.RandomResizedCrop(input_size, scale=(0.75,1.0), ratio=(0.75,1.0), interpolation=PIL.Image.NEAREST), \ transforms.RandomHorizontalFlip()]) train_dataset = CityscapesData(root_img_folder, root_label_folder, train_anno_txt, train_hdf5_file, input_size, cls_num=args.cls_num, img_transform=transforms.Compose([ train_augmentation, RGB2BGR(roll=True), ToTorchFormatTensor(div=False), normalize, ]), label_transform=transforms.Compose([ transforms.ToPILImage(), train_label_augmentation, transforms.ToTensor(), ])) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_dataset = CityscapesData( root_img_folder, root_label_folder, val_anno_txt, val_hdf5_file, input_size, cls_num=args.cls_num, img_transform=transforms.Compose([ transforms.Resize([input_size, input_size]), RGB2BGR(roll=True), ToTorchFormatTensor(div=False), normalize, ]), label_transform=transforms.Compose([ transforms.ToPILImage(), transforms.Resize([input_size, input_size], interpolation=PIL.Image.NEAREST), transforms.ToTensor(), ])) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size / 2, shuffle=False, num_workers=args.workers, pin_memory=True) return train_loader, val_loader
# Initialize the model for this run model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True) # Print the model we just instantiated print(model_ft) # -------------------------- LOADING THE DATA -------------------------- # Data augmentation and normalization for training # Just normalization for validation data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(input_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(input_size), transforms.CenterCrop(input_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } print("Initializing Datasets and Dataloaders...")
import torch import torch.nn.functional as F from torch import nn from torch import optim from torchvision import datasets, transforms, models from collections import OrderedDict import json data_dir = 'flowers' train_dir = data_dir + '/train' valid_dir = data_dir + '/valid' test_dir = data_dir + '/test' device = "cuda" transform_train = transforms.Compose([transforms.RandomResizedCrop(224), transforms.RandomRotation(30), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) transform_test = transforms.Compose([transforms.Resize(256), transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) transform_valid = transform_test trainset = datasets.ImageFolder(train_dir, transform = transform_train) testset = datasets.ImageFolder(test_dir, transform = transform_test) validset = datasets.ImageFolder(valid_dir, transform = transform_valid)
def main_worker(gpu, ngpus_per_node, args): 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) logger = setup_logger(output='output/{}/'.format(args.output_dir), distributed_rank=dist.get_rank(), name="moco") # create model logger.info("=> creating model '{}'".format(args.arch)) model = moco.builder.MoCo( models.__dict__[args.arch], args.moco_dim, args.moco_k, args.moco_m, args.moco_t, args.mlp) logger.info(model) 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): logger.info("=> 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']) logger.info("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: logger.info("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code traindir = os.path.join(args.data, 'train') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if args.aug_plus: # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709 augmentation = [ transforms.RandomResizedCrop(224, scale=(0.2, 1.)), 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([.1, 2.])], p=0.5), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ] else: # MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978 augmentation = [ transforms.RandomResizedCrop(224, scale=(0.2, 1.)), transforms.RandomGrayscale(p=0.2), transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ] train_dataset = datasets.ImageFolder( traindir, moco.loader.TwoCropsTransform(transforms.Compose(augmentation))) 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, drop_last=True) 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, logger, args) 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(), 'optimizer' : optimizer.state_dict(), }, is_best=False, filename='output/{}/checkpoint_current.pth.tar'.format(args.output_dir)) if epoch % args.save_freq == 0: save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'optimizer' : optimizer.state_dict(), }, is_best=False, filename='output/{}/checkpoint_{:04d}.pth.tar'.format(args.output_dir,epoch))
def make_dataset(): # Small noise is added, following SN-GAN def noise(x): return x + torch.FloatTensor(x.size()).uniform_(0, 1.0 / 128) if opt.dataset == "cifar10": trans = tfs.Compose([ tfs.RandomCrop(opt.img_width, padding=4), tfs.RandomHorizontalFlip(), tfs.ToTensor(), tfs.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]), tfs.Lambda(noise) ]) data = CIFAR10(root=opt.root, train=True, download=True, transform=trans) loader = DataLoader(data, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers) elif opt.dataset == "dog_and_cat_64": trans = tfs.Compose([ tfs.RandomResizedCrop(opt.img_width, scale=(0.8, 0.9), ratio=(1.0, 1.0)), tfs.RandomHorizontalFlip(), tfs.ToTensor(), tfs.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]), tfs.Lambda(noise) ]) data = ImageFolder(opt.root, transform=trans) loader = DataLoader(data, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers) elif opt.dataset == "dog_and_cat_128": trans = tfs.Compose([ tfs.RandomResizedCrop(opt.img_width, scale=(0.8, 0.9), ratio=(1.0, 1.0)), tfs.RandomHorizontalFlip(), tfs.ToTensor(), tfs.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]), tfs.Lambda(noise) ]) data = ImageFolder(opt.root, transform=trans) loader = DataLoader(data, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers) elif opt.dataset == "imagenet": trans = tfs.Compose([ tfs.RandomResizedCrop(opt.img_width, scale=(0.8, 0.9), ratio=(1.0, 1.0)), tfs.RandomHorizontalFlip(), tfs.ToTensor(), tfs.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]), tfs.Lambda(noise) ]) data = ImageFolder(opt.root, transform=trans) loader = DataLoader(data, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers) else: raise ValueError(f"Unknown dataset: {opt.dataset}") return loader
SIZE = 224 transform_test = transforms.Compose([ transforms.Resize(SIZE), transforms.CenterCrop(SIZE), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) transform_train = transforms.Compose([ transforms.ColorJitter(brightness=0.4, contrast=0.2, saturation=0.2, hue=0.2), transforms.RandomRotation(15), transforms.RandomResizedCrop(SIZE, scale=(0.8, 1.0), ratio=(0.90, 1.10)), transform_test, ]) def imshow(img): img = img / 2 + 0.5 img = img.numpy() plt.imshow(img.transpose((1, 2, 0))) plt.show() def imload(path): return Image.open(path).convert("RGB")
def main(): ia.seed(1) train_datapath = "./food11re/skewed_training" valid_datapath = "./food11re/validation" test_datapath = "./food11re/evaluation" transform = transforms.Compose([ transforms.RandomRotation(30), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), ImgAugTransform(), lambda x: PIL.Image.fromarray(x), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_dataset = Food11Dataset(train_datapath, is_train=True) train_dataset_folder = torchvision.datasets.ImageFolder( root='./food11re/skewed_training', transform=transform) valid_dataset = Food11Dataset(valid_datapath, is_train=False) test_dataset = Food11Dataset(test_datapath, is_train=False) #wts = [100, 781, 67, 169, 196, 75, 757, 1190, 194, 67, 2857] #train_dataset.augmentation(wts) weight = [] for i in range(11): class_count = train_dataset_folder.targets.count(i) weight.append(1. / (class_count / len(train_dataset_folder.targets))) samples_weight = np.array([weight[t] for _, t in train_dataset_folder]) weighted_sampler = data.WeightedRandomSampler(samples_weight, num_samples=15000, replacement=True) randon_sampler = data.RandomSampler(train_dataset, replacement=True, num_samples=9000, generator=None) print( "----------------------------------------------------------------------------------" ) print("Dataset bf. loading - ", train_datapath) print(train_dataset.show_details()) print( "----------------------------------------------------------------------------------" ) print("Dataset bf. loading - ", valid_datapath) print(valid_dataset.show_details()) print( "----------------------------------------------------------------------------------" ) print("Dataset bf. loading - ", test_datapath) print(test_dataset.show_details()) train_folder_loader = DataLoader(dataset=train_dataset_folder, num_workers=0, batch_size=100, sampler=weighted_sampler) train_loader = DataLoader(dataset=train_dataset, num_workers=0, batch_size=100, sampler=randon_sampler) valid_loader = DataLoader(dataset=valid_dataset, num_workers=0, batch_size=100, shuffle=False) test_loader = DataLoader(dataset=test_dataset, num_workers=0, batch_size=100, shuffle=False) data_loading(train_folder_loader, train_dataset) data_loading(train_loader, train_dataset) data_loading(valid_loader, valid_dataset) data_loading(test_loader, test_dataset)
import torch.optim as optim import pickle # In[73]: num_classes = 120 batch_size = 13 epochs = 15 sample_submission = pd.read_csv('./data/sample_submission.csv') # In[74]: normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) train_transform = transforms.Compose([ transforms.RandomResizedCrop(299), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ]) valid_transform = transforms.Compose([ transforms.Resize(360), transforms.CenterCrop(299), transforms.ToTensor(), normalize ]) trainset = torchvision.datasets.ImageFolder("./data/train/", train_transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)