def main(batch_size, saves_dir=TENSORFLOW_SAVES_DIR): saves_dir = os.path.join(BASE_DIR, saves_dir) batches = [1, 8, 16, 32, 64] if batch_size: batches = [batch_size] optimized_graph = os.path.join(saves_dir, 'optimized_graph.pb') print("Optimized graph size:", os.path.getsize(optimized_graph)) quantized_graph = os.path.join(saves_dir, 'quantized_graph.pb') print("Quantized graph size:", os.path.getsize(quantized_graph)) for batch_size in batches: print("Batch size: {}".format(batch_size)) batch = np.random.random((batch_size, 224, 224, 3)) tf.reset_default_graph() frozen_model = OptimizedModel(saves_dir, input_node_name=Model.input_node_name, output_node_name=Model.output_node_name, graph_name='optimized_graph.pb') measure_model(frozen_model, "Optim. model", batch) frozen_model.sess.close() tf.reset_default_graph() frozen_model = OptimizedModel(saves_dir, input_node_name=Model.input_node_name, output_node_name=Model.output_node_name, graph_name='quantized_graph.pb') measure_model(frozen_model, "Quantized model", batch) frozen_model.sess.close()
def main(batch_size, saves_dir=TENSORFLOW_SAVES_DIR): saves_dir = os.path.join(BASE_DIR, saves_dir) batches = [1, 8, 16, 32, 64] if batch_size: batches = [batch_size] for batch_size in batches: print("Batch size: {}".format(batch_size)) batch = np.random.random((batch_size, 224, 224, 3)) tf.reset_default_graph() usual_model = Model() measure_model(usual_model, "Usual model", batch) usual_model.sess.close() tf.reset_default_graph() frozen_model = FrozenModel(saves_dir) measure_model(frozen_model, "Frozen model", batch) frozen_model.sess.close()
def evaluate(): # build dataset val_loader, n_class = get_dataset() # build model net = get_model(n_class) # for measure n_flops, n_params = measure_model(net, 224, 224) print('=> Model Parameter: {:.3f} M, FLOPs: {:.3f}M'.format(n_params / 1e6, n_flops / 1e6)) del net net = get_model(n_class) criterion = nn.CrossEntropyLoss() if use_cuda: net = net.cuda() net = torch.nn.DataParallel(net, list(range(args.n_gpu))) cudnn.benchmark = True # begin eval net.eval() batch_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() end = time.time() with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(val_loader): if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() inputs, targets = Variable(inputs), Variable(targets) outputs = net(inputs) loss = criterion(outputs, targets) # measure accuracy and record loss prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5)) losses.update(loss.item(), inputs.size(0)) top1.update(prec1.item(), inputs.size(0)) top5.update(prec5.item(), inputs.size(0)) # timing batch_time.update(time.time() - end) end = time.time() progress_bar(batch_idx, len(val_loader), 'Loss: {:.3f} | Acc1: {:.3f}% | Acc5: {:.3f}%' .format(losses.avg, top1.avg, top5.avg))
def main(): # set model model = getattr(models, args.model)(args) if args.data == 'cifar10': image_size = 32 args.num_classes = 10 elif args.data == 'cifar100': image_size = 32 args.num_classes = 100 elif args.data == 'imagenet': image_size = 224 args.num_classes = 1000 else: raise NotImplementedError n_flops, n_params = measure_model(model, image_size, image_size) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) if torch.cuda.device_count(): model = torch.nn.DataParallel(model) # for multi-GPU training if torch.cuda.is_available(): model.cuda() print(model) if args.mode == 'train': # get the training loader and validation loader train_set, val_set = read_train_data(datadir=args.data_dir, data=args.data) # set the start epoch value if args.resume: start_epoch = None else: start_epoch = args.start_epoch train(startepoch=start_epoch, epochs=args.epochs, model=model, train_set=train_set, val_set=val_set, resume=args.resume) elif args.mode == 'test': test_set = read_test_data(datadir=args.data_dir, data=args.data, mode='test') test(model=model, test_set=test_set) else: raise NotImplementedError
def main(batch_size): batches = [1, 8, 16, 32, 64] if batch_size: batches = [batch_size] pt_model = PTModel() keras_model = KerasModel() tf_model = TFModel() for batch_size in batches: print("Batch size: {}".format(batch_size)) pt_batch = np.random.random((batch_size, 3, 224, 224)) measure_model(pt_model, "Pytorch", pt_batch) tf_batch = np.random.random((batch_size, 224, 224, 3)) measure_model(keras_model, "Keras", tf_batch) measure_model(tf_model, "Tensorflow", tf_batch)
nn.LeakyReLU(0.2, inplace=True) ]) self.block4 = nn.Sequential(*[ nn.Conv2d(256, 512, 4, padding=1), nn.InstanceNorm2d(512, affine=True), nn.LeakyReLU(0.2, inplace=True) ]) # FCN classification layer self.fc = nn.Conv2d(512, 1, 4, padding=1) def forward(self, x): h1 = self.block1(x) h2 = self.block2(h1) h3 = self.block3(h2) h4 = self.block4(h3) out = self.fc(h4) # Average pooling and flatten out = F.avg_pool2d(out, out.size()[2:]).view(out.size()[0], -1) D_outputs = namedtuple("DOutputs", ['h1', 'h2', 'h3', 'h4', 'out']) return D_outputs(h1, h2, h3, h4, out) if __name__ == "__main__": from utils import measure_model, model_size g = Generator(3, 3, transconv=False) measure_model(g, 256, 256) print(model_size(g))
def main(**kwargs): global args, best_prec1 # Override if needed for arg, v in kwargs.items(): args.__setattr__(arg, v) ### Calculate FLOPs & Param model = getattr(models, args.model)(args) if args.data in ['cifar10', 'cifar100']: IMAGE_SIZE = 32 else: IMAGE_SIZE = 224 n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE, args.debug) if 'measure_only' in args and args.measure_only: return n_flops, n_params print('Starting.. FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) args.filename = "%s_%s_%s.txt" % \ (args.model, int(n_params), int(n_flops)) del (model) # Create model model = getattr(models, args.model)(args) if args.debug: print(args) print(model) model = torch.nn.DataParallel(model).cuda() # Define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Resume from a checkpoint if args.resume: checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) # Evaluate from a model if args.evaluate_from is not None: args.evaluate = True state_dict = torch.load(args.evaluate_from)['state_dict'] model.load_state_dict(state_dict) cudnn.benchmark = True ### Data loading if args.data == "cifar10": train_set = datasets.CIFAR10('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ])) val_set = datasets.CIFAR10('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), ])) elif args.data == "cifar100": train_set = datasets.CIFAR100('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ])) val_set = datasets.CIFAR100('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), ])) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) # Run Forward / Backward passes if args.evaluate: validate(val_loader, model, criterion) return for epoch in range(args.start_epoch, args.epochs): # Train for one epoch tr_prec1, tr_prec5, loss, lr = \ train(train_loader, model, criterion, optimizer, epoch) # Evaluate on validation set val_prec1, val_prec5 = validate(val_loader, model, criterion) # Remember best prec@1 and save checkpoint is_best = val_prec1 < best_prec1 best_prec1 = max(val_prec1, best_prec1) model_filename = 'checkpoint_%03d.pth.tar' % epoch save_checkpoint( { 'epoch': epoch, 'model': args.model, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, args, is_best, model_filename, "%.4f %.4f %.4f %.4f %.4f %.4f\n" % (val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr)) # TestModel and return model = model.cpu().module model = nn.DataParallel(model).cuda() print(model) validate(val_loader, model, criterion) n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('Finished training! FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) print('Please run again with --resume --evaluate flags,' ' to evaluate the best model.') return
def main(**kwargs): global args, best_prec1 mapper = { 1 : 'airplane', 2 : 'automobile', 3 : 'bird', 4 : 'cat', 5 : 'deer', 6 : 'dog', 7 : 'frog', 8 : 'horse', 9 : 'ship', 10 : 'truck' } # Override if needed for arg, v in kwargs.items(): args.__setattr__(arg, v) ### Calculate FLOPs & Param model = getattr(models, args.model)(args) if args.data in ['cifar10', 'cifar100']: IMAGE_SIZE = 32 else: IMAGE_SIZE = 224 n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE, args.debug) if 'measure_only' in args and args.measure_only: return n_flops, n_params print('Starting.. FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) args.filename = "%s_%s_%s.txt" % \ (args.model, int(n_params), int(n_flops)) del(model) # Create model model = getattr(models, args.model)(args) if args.debug: print(args) print(model) criterion = nn.CrossEntropyLoss() if torch.cuda.is_available(): model = torch.nn.DataParallel(model).cuda() criterion = criterion.cuda() # Define loss function (criterion) and optimizer optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) checkpoint = load_checkpoint(args) args.start_epoch = checkpoint['epoch'] + 1 best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) cudnn.benchmark = True ### Data loading if args.data == "cifar10": train_set = datasets.CIFAR10('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ])) val_set = datasets.CIFAR10('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), ])) elif args.data == "cifar100": train_set = datasets.CIFAR100('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ])) val_set = datasets.CIFAR100('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), ])) val_loader = torch.utils.data.DataLoader( val_set, batch_size=2, shuffle=False, num_workers=args.workers, pin_memory=True) batch_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() try: top1_per_cls = [AverageMeter() for i in range(0, model.module.num_blocks)] top5_per_cls = [AverageMeter() for i in range(0, model.module.num_blocks)] except: top1_per_cls = [AverageMeter() for i in range(0, model.num_blocks)] top5_per_cls = [AverageMeter() for i in range(0, model.num_blocks)] model.eval() end = time.time() from gradCam import * for i, (input, target) in enumerate(val_loader): if args.imgNo > 0 and i!=args.imgNo: continue (inputM, targetM) = (input, target) input, target = input[0], target[0].view(1) input = input.view(1,input.shape[0], input.shape[1],input.shape[2]) if torch.cuda.is_available(): target = target.cuda(async=True) input_var = torch.autograd.Variable(input, requires_grad=True) target_var = torch.autograd.Variable(target) # ### Compute output scores, feats = model(input_var, 0.0, p=1) if args.model == 'msdnet': loss = msd_loss(scores, target_var, criterion) else: loss = criterion(scores, target_var) ############################################################################# name = 'diags/' + mapper[target_var.cpu().data.numpy()[0]+1] + '_' + str(args.imgNo) + '_' + str(args.maxC) name2 = 'diags/' + mapper[target_var.cpu().data.numpy()[0]+1] + '_' + str(args.imgNo) if len(args.classLabel): if mapper[target_var.cpu().data.numpy()[0]+1] == args.classLabel: print(correctNess(scores, target_var), i, mapper[target_var.cpu().data.numpy()[0]+1]) elif args.imgNo < 0: print(correctNess(scores, target_var), i, mapper[target_var.cpu().data.numpy()[0]+1]) elif i == args.imgNo: import os if not os.path.exists('diags'): os.makedirs('diags') print("Category : {}".format(mapper[target_var.cpu().data.numpy()[0]+1])) print(correctNess(scores, target_var)) grad_cam = GradCam(model = model) mask = grad_cam(target_var.cpu().data.numpy()[0], input_var, 0, args.maxC-1)#target_index) gb_model = GuidedBackpropReLUModel(model) gb = gb_model(target_var.cpu().data.numpy()[0], input_var, 0, args.maxC-1).transpose(2,0,1) img = input[0].cpu().data.numpy().transpose(1,2,0) img = cv2.resize(img, (512, 512)) show_cam_on_image(img, mask, name) utils.save_image(torch.from_numpy(gb), name +'_gb.jpg') cam_mask = np.zeros(gb.shape) for i in range(0, gb.shape[0]): cam_mask[i, :, :] = mask cam_gb = np.multiply(cam_mask, gb) utils.save_image(torch.from_numpy(cam_gb), name +'_cam_gb.jpg') img = cv2.resize(input.cpu().data.numpy()[0].transpose(1,2,0), (512, 512)) utils.save_image(torch.from_numpy(img.transpose(2,0,1)), name2 + '_input.jpg') exit() else: continue
def main(): global args, model ### Calculate FLOPs & Param model = getattr(models, args.model)(args) IMAGE_SIZE = args.inp_img_size n_flops, n_params = measure_model(model.cuda(), IMAGE_SIZE, IMAGE_SIZE, args.scaling_factor) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) args.filename = "%s_%s_%s.txt" % \ (args.model, int(n_params), int(n_flops)) del (model) ########### Building Model ############## cudnn.benchmark = True model = getattr(models, args.model)(args) model = torch.nn.DataParallel(model).cuda() criterion = charbonnier_loss().cuda() optimizer = torch.optim.Adam(model.parameters(), args.lr) lr = args.lr ########## Data Loading ############### traindir = os.path.join("../data/", args.train_data) testdir = os.path.join("../data/", args.test_data) train_img = read_data(traindir, crop_size=args.inp_img_size, upscale_factor=args.scaling_factor, c_dim=args.c_dim, stride=128) train_set = TensorClass(train_img[0], train_img[1], train_img[2]) test_set = testDatasetFromFolder2(testdir, upscale_factor=args.scaling_factor) train_loader = data.DataLoader(dataset=train_set, num_workers=args.workers, batch_size=args.batch_size, shuffle=True) test_loader = data.DataLoader(dataset=test_set, num_workers=args.workers, batch_size=test_batch_size, shuffle=False) ### Optionally resume from a checkpoint if args.resume: checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) ### Optionally evaluate from a model if args.evaluate_from is not None: args.evaluate = True state_dict = torch.load(args.evaluate_from)['state_dict'] model.load_state_dict(state_dict) if args.evaluate: test_psnr, test_ssim = test(test_loader, model, criterion) print("test Psnr: " + str(test_psnr) + "\t test SSIM: " + str(test_ssim)) return for epoch in range(args.start_epoch, args.epochs): ### Train for one epoch tr_psnr, tr_ssim ,loss, lr = \ train(train_loader, model, criterion, optimizer, epoch) ### Evaluate on test set test_psnr, test_ssim = test(test_loader, model, criterion) print("Train Psnr: " + str(tr_psnr) + "\t Train SSIM: " + str(tr_ssim) + "\t Loss: " + str(loss)) print("test Psnr: " + str(test_psnr) + "\t test SSIM: " + str(test_ssim)) ### save checkpoint model_filename = 'checkpoint_%03d.pth.tar' % epoch save_checkpoint( { 'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'psnr': best_psnr, }, args, model_filename)
def main(): global args, best_acc1 args.filename = 'log.txt' if args.dataset in ['cifar10', 'cifar100']: IMAGE_SIZE = 32 else: IMAGE_SIZE = 224 fd = open(args.record_file, 'a') print('Args Config:', str(args)) fd.write(str(args) + '\n') if args.continue_train: print("Continue training!") fd.write("Continue training!" + "\n") args.resume = args.train_url + "save_models/checkpoint.pth.tar" ### Calculate FLOPs & Param else: if args.model.startswith("G"): model = GhostNet.get_GhostNet(args) else: model = MicroNet.get_MicroNet(args) n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) fd.write('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6) + '\n') print('Model Struture:', str(model)) fd.write(str(model) + '\n') del (model) fd.close() # model = MicroNet.get_MicroNet(args) if args.model.startswith("G"): model = GhostNet.get_GhostNet(args) else: model = MicroNet.get_MicroNet(args) model.cuda() ### Define loss function (criterion) and optimizer if args.LSR_Mixup: criterion = CrossEntryLoss_onehot().cuda() else: criterion = nn.CrossEntropyLoss().cuda() criterion_temper = nn.CrossEntropyLoss().cuda() # for validation optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) ### Optionally resume from a checkpoint if args.resume: checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 best_acc1 = checkpoint['best_acc1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) cudnn.benchmark = True if args.train_on_cloud: args.data_url = "./data" ### Data loading if args.dataset == "cifar10": normalize = transforms.Normalize(mean=[0.4914, 0.4824, 0.4467], std=[0.2471, 0.2435, 0.2616]) train_set = datasets.CIFAR10(args.data_url, train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4, padding_mode=args.pad_mode), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.CIFAR10(args.data_url, train=False, transform=transforms.Compose([ transforms.ToTensor(), normalize, ])) elif args.dataset == "cifar100": normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408], std=[0.2675, 0.2565, 0.2761]) train_set = datasets.CIFAR100(args.data_url, train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4, padding_mode=args.pad_mode), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.CIFAR100(args.data_url, train=False, transform=transforms.Compose([ transforms.ToTensor(), normalize, ])) elif args.dataset == "mini-imagenet": normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_set = mini_imagenet(args.data_url, "split_train.csv", transforms.Compose([ transforms.RandomSizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = mini_imagenet(args.data_url, "split_val.csv", transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])) else: traindir = os.path.join(args.data_url, 'train') valdir = os.path.join(args.data_url, 'val') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_set = datasets.ImageFolder(traindir, transforms.Compose([ transforms.RandomSizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.ImageFolder(valdir, transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])) train_loader = torch.utils.data.DataLoader( train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False) val_loader = torch.utils.data.DataLoader( val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False) if args.evaluate: criterion_temper = nn.CrossEntropyLoss.cuda() validate(val_loader, model, criterion_temper) # TODO: validate must before measure_model otherwise result be worese, but I don't know why n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) fd = open(args.record_file, 'a') print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) fd.write('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6) + '\n') return epoch_time = AverageMeter() start_time = time.time() for epoch in range(args.start_epoch, args.epochs): ### Train for one epoch tr_acc1, tr_acc5, tr_loss, lr = \ train(train_loader, model, criterion, optimizer, epoch, args) ### Evaluate on validation set val_acc1, val_acc5, val_loss = validate(val_loader, model, criterion_temper) ### Remember best Acc@1 and save checkpoint is_best = val_acc1 > best_acc1 best_acc1 = max(val_acc1, best_acc1) # model_filename = 'checkpoint_%03d.pth.tar' % epoch model_filename = 'checkpoint.pth.tar' save_checkpoint({ 'epoch': epoch, 'model': args.model, 'state_dict': model.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), }, args, is_best, model_filename, "%3d %.4f %.4f %.4f %.4f %.4f %.4f %.4f\n" % (epoch, val_acc1, val_acc5, tr_acc1, tr_acc5, val_loss, tr_loss, lr)) epoch_time.update(time.time() - start_time, 1) string = 'Avg_time: %4f H, Duration: %4f H, Left Time: %4f H' % \ (epoch_time.avg / 3600, epoch_time.sum / 3600, epoch_time.avg * (args.epochs - epoch - 1) / 3600) print(string) fd = open(args.record_file, 'a') fd.write(string + '\n') fd.close() start_time = time.time() ### Convert model and test # model = model.cpu().module # print('Start Converting ...') # convert_model(model, args) # print('Converting End!') print('Model Struture:', str(model)) validate(val_loader, model, criterion_temper) model.cpu() n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) fd = open(args.record_file, 'a') fd.write(str(model) + '\n') fd.write('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6) + '\n') fd.close() return
def main(): global args, best_prec1 ### Calculate FLOPs & Param model = getattr(models, args.model)(args) print(model) if args.data in ['cifar10', 'cifar100']: IMAGE_SIZE = 32 else: IMAGE_SIZE = 224 n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) del (model) args.no_save_model = True print(args) ### Create model model = getattr(models, args.model)(args) if args.model.startswith('alexnet') or args.model.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() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) ### Optionally resume from a checkpoint if args.resume: checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) ### Optionally convert from a model and saves it to onnx file if args.convert_from is not None: args.evaluate = True state_dict = torch.load(args.convert_from)['state_dict'] model.load_state_dict(state_dict) model = model.cpu().module convert_model(model, args) model = model.cuda() # head, tail = os.path.split(args.convert_from) # tail = "converted_" + tail # torch.save({'state_dict': model.state_dict()}, os.path.join(head, tail)) dummy_input = torch.randn(args.batch_size, 3, IMAGE_SIZE, IMAGE_SIZE, device='cuda') torch.onnx.export(model, dummy_input, args.convert_from + ".onnx", verbose=True) return ### Optionally evaluate from a model if args.evaluate_from is not None: args.evaluate = True state_dict = torch.load(args.evaluate_from)['state_dict'] model.load_state_dict(state_dict) cudnn.benchmark = True ### Data loading if args.data == "cifar10": norm_mean = [0.49139968, 0.48215827, 0.44653124] norm_std = [0.24703233, 0.24348505, 0.26158768] norm_transform = transforms.Normalize(norm_mean, norm_std) train_set = datasets.CIFAR10('~/Documents/CIFAR-10', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop( 32, padding=4, padding_mode='reflect'), transforms.RandomHorizontalFlip(), transforms.ToTensor(), norm_transform, ])) val_set = datasets.CIFAR10('~/Documents/CIFAR-10', train=False, transform=transforms.Compose([ transforms.ToTensor(), norm_transform, ])) elif args.data == "cifar100": normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408], std=[0.2675, 0.2565, 0.2761]) train_set = datasets.CIFAR100('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.CIFAR100('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), normalize, ])) else: 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_set = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomSizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.ImageFolder( valdir, transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: validate(val_loader, model, criterion) return for epoch in range(args.start_epoch, args.epochs): ### Train for one epoch tr_prec1, tr_prec5, loss, lr = \ train(train_loader, model, criterion, optimizer, epoch) ### Evaluate on validation set val_prec1, val_prec5 = validate(val_loader, model, criterion) ### Remember best prec@1 is_best = val_prec1 < best_prec1 best_prec1 = max(val_prec1, best_prec1) save_checkpoint( { 'epoch': epoch, 'model': args.model, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, args, is_best, args.file_name) ### Convert model and test model = model.cpu().module convert_model(model, args) model = nn.DataParallel(model).cuda() print(model) validate(val_loader, model, criterion) n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) return
def main(**kwargs): global args, best_prec1, acc_file_name, flops_file_name # Override if needed # kwargs is a dict for arg, v in kwargs.items(): args.__setattr__(arg, v) # set args.arg=v, copy kwargs to args ### Calculate FLOPs & Param model = getattr(models, args.model)(args) ###### ? get models.(args.model) if args.data in ['cifar10', 'cifar100']: IMAGE_SIZE = 32 if (args.data == 'cifar10'): acc_file_name = "cifar10_acc.csv" flops_file_name = "cifar10_flops.csv" else: acc_file_name = "cifar100_acc.csv" flops_file_name = "cifar100_flops.csv" else: IMAGE_SIZE = 224 set_save_path() # calculate the FLOPs # if (args.evaluate): n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE, flops_file_path, args.debug) ##### if 'measure_only' in args and args.measure_only: # no 'measure_only' parameter return n_flops, n_params print('Starting.. FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) args.filename = "%s_%s_%s.txt" % (args.model, int(n_params), int(n_flops)) # del(model) # Create model(the 2nd) # model = getattr(models, args.model)(args) #### if args.debug: print(args) # print(model) model = torch.nn.DataParallel( model).cuda() # Implements data parallelism at the module level # if exist mulit-devices, This container parallelizes the application of the given module by splitting # the input across the specified devices by chunking in the batch dimension (other objects will be copied once per device) ### Data loading, no nomarlisation if args.data == "cifar10": train_set = datasets.CIFAR10('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ])) val_set = datasets.CIFAR10('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), ])) elif args.data == "cifar100": train_set = datasets.CIFAR100('../data', train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ])) val_set = datasets.CIFAR100('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), ])) cudnn.benchmark = True ###### train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) # test_loader = torch.utils.data.DataLoader( # val_set, # batch_size=args.batch_size, shuffle=False, # num_workers=args.workers, pin_memory=True) test_loader = val_loader # Define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) ##### validate # Resume from a checkpoint if args.resume: # use latest checkpoint if have any checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) # Evaluate from a model if args.evaluate_from is not None: # path to saved checkpoint args.evaluate = True state_dict = torch.load(args.evaluate_from)['state_dict'] model.load_state_dict(state_dict) if args.evalmode is not None: if args.evalmode == 'anytime': validate(val_loader, model, criterion) else: print('dynamic') dynamic_evaluate(model, test_loader, val_loader, args) return # Run Forward / Backward passes # evaluate and return # if args.evaluate: # validate(val_loader, model, criterion) # return ###### train best_epoch = 0 for epoch in range(args.start_epoch, args.epochs): # Train for one epoch tr_prec1, tr_prec5, loss, lr = train(train_loader, model, criterion, optimizer, epoch) # Evaluate on validation set val_prec1, val_prec5 = validate(val_loader, model, criterion) # Remember best prec@1 and save checkpoint is_best = val_prec1 < best_prec1 best_prec1 = max(val_prec1, best_prec1) if (is_best): best_epoch = epoch model_filename = 'checkpoint_%03d.pth.tar' % epoch save_checkpoint( { 'epoch': epoch, 'model': args.model, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, args, is_best, model_filename, "%.4f %.4f %.4f %.4f %.4f %.4f\n" % (val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr)) print('Best val_prec1: {:.4f} at epoch {}'.format(best_prec1, best_epoch)) # TestModel and return model = model.cpu().module model = nn.DataParallel(model).cuda() # if args.debug: # print(model) validate(val_loader, model, criterion) # n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE, flops_file_path,args.debug) # print('Finished training! FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) print('Please run again with --resume --evaluate flags,' ' to evaluate the best model.') return
args = parser.parse_args() # ------------------ MEAN & STD --------------------- PIXEL_MEANS = np.array([0.485, 0.456, 0.406]) PIXEL_STDS = np.array([0.229, 0.224, 0.225]) # --------------------------------------------------- # Set GPU id, CUDA and cudnn os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id USE_CUDA = torch.cuda.is_available() cudnn.benchmark = True # Create & Load model MODEL = models.__dict__[args.arch]() # Calculate FLOPs & Param n_flops, n_convops, n_params = measure_model(MODEL, args.crop_size, args.crop_size) print('==> FLOPs: {:.4f}M, Conv_FLOPs: {:.4f}M, Params: {:.4f}M'. format(n_flops / 1e6, n_convops / 1e6, n_params / 1e6)) del MODEL # Load Weights MODEL = models.__dict__[args.arch]() checkpoint = torch.load(args.model_weights) weight_dict = checkpoint model_dict = MODEL.state_dict() updated_dict, match_layers, mismatch_layers = weight_filler(weight_dict, model_dict) model_dict.update(updated_dict) MODEL.load_state_dict(model_dict) # Switch to evaluate mode MODEL.cuda().eval()
def main(): global args, best_prec1 ### Calculate FLOPs & Param model = getattr(models, args.model)(args) print(model) if args.data in ['cifar10', 'cifar100']: IMAGE_SIZE = 32 else: IMAGE_SIZE = 224 n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) args.filename = "%s_%s_%s.txt" % \ (args.model, int(n_params), int(n_flops)) del (model) print(args) torch.manual_seed(args.manual_seed) # torch.cuda.manual_seed_all(args.manual_seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False np.random.seed(args.manual_seed) ### Create model model = getattr(models, args.model)(args) if args.model.startswith('alexnet') or args.model.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() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) ### Optionally resume from a checkpoint if args.resume: checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) ### Optionally convert from a model if args.convert_from is not None: args.evaluate = True state_dict = torch.load(args.convert_from)['state_dict'] model.load_state_dict(state_dict) model = model.cpu().module convert_model(model, args) model = nn.DataParallel(model).cuda() head, tail = os.path.split(args.convert_from) tail = "converted_" + tail torch.save({'state_dict': model.state_dict()}, os.path.join(head, tail)) ### Optionally evaluate from a model if args.evaluate_from is not None: args.evaluate = True state_dict = torch.load(args.evaluate_from)['state_dict'] model.load_state_dict(state_dict) ### Data loading dataset_dir = args.dataset_dir if args.dataset_dir is not None else '../data' if args.data == "cifar10": normalize = transforms.Normalize(mean=[0.4914, 0.4824, 0.4467], std=[0.2471, 0.2435, 0.2616]) train_set = datasets.CIFAR10(dataset_dir, train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.CIFAR10(dataset_dir, train=False, transform=transforms.Compose([ transforms.ToTensor(), normalize, ])) elif args.data == "cifar100": normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408], std=[0.2675, 0.2565, 0.2761]) train_set = datasets.CIFAR100(dataset_dir, train=True, download=True, transform=transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.CIFAR100(dataset_dir, train=False, transform=transforms.Compose([ transforms.ToTensor(), normalize, ])) else: 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_set = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomSizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) val_set = datasets.ImageFolder( valdir, transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: validate(val_loader, model, criterion) return for epoch in range(args.start_epoch, args.epochs): ### Train for one epoch tr_prec1, tr_prec5, loss, lr = \ train(train_loader, model, criterion, optimizer, epoch) ### Evaluate on validation set val_prec1, val_prec5 = validate(val_loader, model, criterion) ### Remember best prec@1 and save checkpoint is_best = val_prec1 < best_prec1 best_prec1 = max(val_prec1, best_prec1) model_filename = 'checkpoint_%03d.pth.tar' % epoch save_checkpoint( { 'epoch': epoch, 'model': args.model, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, args, is_best, model_filename, "%.4f %.4f %.4f %.4f %.4f %.4f\n" % (val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr)) ### Convert model and test model = model.cpu().module convert_model(model, args) model = nn.DataParallel(model).cuda() print(model) validate(val_loader, model, criterion) n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) return
def run_training(args): # create model model = models.__dict__[args.arch](args.pretrained) model = torch.nn.DataParallel(model).cuda() # measure FLOPs: measure_model(model, 32, 32) best_prec1 = 0 # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): logging.info('=> loading checkpoint `{}`'.format(args.resume)) checkpoint = torch.load(args.resume) args.start_iter = checkpoint['iter'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) logging.info('=> loaded checkpoint `{}` (iter: {})'.format( args.resume, checkpoint['iter'] )) else: logging.info('=> no checkpoint found at `{}`'.format(args.resume)) cudnn.benchmark = True train_loader = prepare_train_data(dataset=args.dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers) test_loader = prepare_test_data(dataset=args.dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() end = time.time() if args.defend_algo == 'ifgm': defend_algo = ifgm elif args.defend_algo == 'None': defend_algo = None else: raise NotImplementedError # adversarial algorithm if args.attack_algo == 'ifgm': attack_algo = ifgm elif args.attack_algo == 'None': attack_algo = None else: raise NotImplementedError for i in range(args.start_iter, args.iters): model.train() adjust_learning_rate(args, optimizer, i) input, target = next(iter(train_loader)) # measuring data loading time data_time.update(time.time() - end) target = target.squeeze().long().cuda(async=True) target_var = Variable(target) if defend_algo: input_adv = defend_algo(input, None, F.cross_entropy, y=target_var, eps=args.defend_eps, model=model, steps=args.defend_adv_iter, gamma=args.defend_gamma, randinit=args.defend_randinit).data input_adv_var = Variable(input_adv) input_var = Variable(input) # compute output output = model(input_var) if defend_algo: output_adv = model(input_adv_var) loss = criterion(output, target_var) * 0.5 + criterion(output_adv, target_var) * 0.5 else: loss = criterion(output, target_var) # measure accuracy and record loss prec1, = accuracy(output.data, target, topk=(1,)) losses.update(loss.item(), input.size(0)) top1.update(prec1, input.size(0)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() # print log if i % args.print_freq == 0: logging.info("Iter: [{0}/{1}]\t" "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t" "Data {data_time.val:.3f} ({data_time.avg:.3f})\t" "Loss {loss.val:.3f} ({loss.avg:.3f})\t" "Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t".format( i, args.iters, batch_time=batch_time, data_time=data_time, loss=losses, top1=top1) ) # evaluate every 1000 steps if (i % args.eval_every == 0 and i > 0) or (i == args.iters - 1): prec1 = validate(args, test_loader, model, criterion) if attack_algo: _ = validate_adv(args, test_loader, model, criterion, attack_algo) is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) checkpoint_path = os.path.join(args.save_path, 'checkpoint_{:05d}.pth.tar'.format( i)) save_checkpoint({ 'iter': i, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, }, is_best, filename=checkpoint_path) shutil.copyfile(checkpoint_path, os.path.join(args.save_path, 'checkpoint_latest' '.pth.tar'))
def main(): global BEST_ACC, LR_STATE start_epoch = cfg.CLS.start_epoch # start from epoch 0 or last checkpoint epoch # Create ckpt folder if not os.path.isdir(cfg.CLS.ckpt): mkdir_p(cfg.CLS.ckpt) if args.cfg_file is not None and not cfg.CLS.evaluate: shutil.copyfile( args.cfg_file, os.path.join(cfg.CLS.ckpt, args.cfg_file.split('/')[-1])) # Dataset and Loader normalize = transforms.Normalize(mean=cfg.pixel_mean, std=cfg.pixel_std) train_aug = [ transforms.RandomResizedCrop(cfg.CLS.crop_size), transforms.RandomHorizontalFlip() ] if len(cfg.CLS.rotation) > 0: train_aug.append(transforms.RandomRotation(cfg.CLS.rotation)) if len(cfg.CLS.pixel_jitter) > 0: train_aug.append(RandomPixelJitter(cfg.CLS.pixel_jitter)) if cfg.CLS.grayscale > 0: train_aug.append(transforms.RandomGrayscale(cfg.CLS.grayscale)) train_aug.append(transforms.ToTensor()) train_aug.append(normalize) traindir = os.path.join(cfg.CLS.data_root, cfg.CLS.train_folder) train_loader = torch.utils.data.DataLoader(datasets.ImageFolder( traindir, transforms.Compose(train_aug)), batch_size=cfg.CLS.train_batch, shuffle=True, num_workers=cfg.workers, pin_memory=True) if cfg.CLS.validate or cfg.CLS.evaluate: valdir = os.path.join(cfg.CLS.data_root, cfg.CLS.val_folder) val_loader = torch.utils.data.DataLoader(datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(cfg.CLS.base_size), transforms.CenterCrop(cfg.CLS.crop_size), transforms.ToTensor(), normalize, ])), batch_size=cfg.CLS.test_batch, shuffle=False, num_workers=cfg.workers, pin_memory=True) # Create model model = models.__dict__[cfg.CLS.arch]() print(model) # Calculate FLOPs & Param n_flops, n_convops, n_params = measure_model(model, cfg.CLS.crop_size, cfg.CLS.crop_size) print('==> FLOPs: {:.4f}M, Conv_FLOPs: {:.4f}M, Params: {:.4f}M'.format( n_flops / 1e6, n_convops / 1e6, n_params / 1e6)) del model model = models.__dict__[cfg.CLS.arch]() # Load pre-train model if cfg.CLS.pretrained: print("==> Using pre-trained model '{}'".format(cfg.CLS.pretrained)) pretrained_dict = torch.load(cfg.CLS.pretrained) try: pretrained_dict = pretrained_dict['state_dict'] except: pretrained_dict = pretrained_dict model_dict = model.state_dict() updated_dict, match_layers, mismatch_layers = weight_filler( pretrained_dict, model_dict) model_dict.update(updated_dict) model.load_state_dict(model_dict) else: print("==> Creating model '{}'".format(cfg.CLS.arch)) # Define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() if cfg.CLS.pretrained: def param_filter(param): return param[1] new_params = map( param_filter, filter(lambda p: p[0] in mismatch_layers, model.named_parameters())) base_params = map( param_filter, filter(lambda p: p[0] in match_layers, model.named_parameters())) model_params = [{ 'params': base_params }, { 'params': new_params, 'lr': cfg.CLS.base_lr * 10 }] else: model_params = model.parameters() model = torch.nn.DataParallel(model).cuda() cudnn.benchmark = True optimizer = optim.SGD(model_params, lr=cfg.CLS.base_lr, momentum=cfg.CLS.momentum, weight_decay=cfg.CLS.weight_decay) # Evaluate model if cfg.CLS.evaluate: print('\n==> Evaluation only') test_loss, test_top1, test_top5 = test(val_loader, model, criterion, start_epoch, USE_CUDA) print( '==> Test Loss: {:.8f} | Test_top1: {:.4f}% | Test_top5: {:.4f}%'. format(test_loss, test_top1, test_top5)) return # Resume training title = 'Pytorch-CLS-' + cfg.CLS.arch if cfg.CLS.resume: # Load checkpoint. print("==> Resuming from checkpoint '{}'".format(cfg.CLS.resume)) assert os.path.isfile( cfg.CLS.resume), 'Error: no checkpoint directory found!' checkpoint = torch.load(cfg.CLS.resume) BEST_ACC = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(cfg.CLS.ckpt, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(cfg.CLS.ckpt, 'log.txt'), title=title) logger.set_names([ 'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.' ]) # Train and val for epoch in range(start_epoch, cfg.CLS.epochs): print('\nEpoch: [{}/{}] | LR: {:.8f}'.format(epoch + 1, cfg.CLS.epochs, LR_STATE)) train_loss, train_acc = mixup_train(train_loader, model, criterion, optimizer, epoch, USE_CUDA) if cfg.CLS.validate: test_loss, test_top1, test_top5 = test(val_loader, model, criterion, epoch, USE_CUDA) else: test_loss, test_top1, test_top5 = 0.0, 0.0, 0.0 # Append logger file logger.append([LR_STATE, train_loss, test_loss, train_acc, test_top1]) # Save model save_checkpoint(model, optimizer, test_top1, epoch) # Draw curve try: draw_curve(cfg.CLS.arch, cfg.CLS.ckpt) print('==> Success saving log curve...') except: print('==> Saving log curve error...') logger.close() try: savefig(os.path.join(cfg.CLS.ckpt, 'log.eps')) shutil.copyfile( os.path.join(cfg.CLS.ckpt, 'log.txt'), os.path.join( cfg.CLS.ckpt, 'log{}.txt'.format( datetime.datetime.now().strftime('%Y%m%d%H%M%S')))) except: print('copy log error.') print('==> Training Done!') print('==> Best acc: {:.4f}%'.format(BEST_ACC))
def main(): global args ### Create model and calculate FLOPs & Param and delete model model = ShuffleNetV2(args.symbol) print(model) IMAGE_SIZE = 224 n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) args.filename = "%s_%s_%s.txt" % (args.model, int(n_params), int(n_flops)) print(args) del (model) ### Create model again model = ShuffleNetV2(args.symbol) # print(model) # args.filename = "%s.txt" % (args.model) # print(args) model = torch.nn.DataParallel(model).cuda() ### Define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) ### Optionally resume from a checkpoint if args.resume: checkpoint = load_checkpoint(args) if checkpoint is not None: args.start_epoch = checkpoint['epoch'] + 1 best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) else: best_prec1 = 0 cudnn.benchmark = True ### Data loading normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if not args.evaluate: traindir = os.path.join(args.data, 'train') train_set = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.2, 1.0), ratio=(0.75, 1.33)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ])) valdir = os.path.join(args.data, 'val') val_set = datasets.ImageFolder( valdir, transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize ])) if not args.evaluate: train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) ### Evaluate if args.evaluate: validate(val_loader, model, criterion) return ### Train for epoch in range(args.start_epoch, args.epochs): ### Train for one epoch tr_prec1, tr_prec5, loss, lr = train(train_loader, model, criterion, optimizer, epoch) ### Evaluate on validation set val_prec1, val_prec5 = validate(val_loader, model, criterion) ### Remember best prec@1 and save checkpoint is_best = (val_prec1 > best_prec1) best_prec1 = max(val_prec1, best_prec1) ### Save checkpoint model_filename = 'checkpoint_%03d.pth.tar' % epoch save_checkpoint( { 'epoch': epoch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict() }, args, is_best, model_filename, "%.4f %.4f %.4f %.4f %.4f %.4f\n" % (val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr)) # save_checkpoint({'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}, # args, model_filename, "%.4f %.4f %.4f %.4f\n" % (tr_prec1, tr_prec5, loss, lr)) n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE) print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6)) return