def convert(): args = parser.parse_args() switch_conv_bn_impl('DBB') switch_deploy_flag(False) train_model = build_model(args.arch) if 'hdf5' in args.load: from utils import model_load_hdf5 model_load_hdf5(train_model, args.load) elif os.path.isfile(args.load): print("=> loading checkpoint '{}'".format(args.load)) checkpoint = torch.load(args.load) if 'state_dict' in checkpoint: checkpoint = checkpoint['state_dict'] checkpoint = copy.deepcopy(checkpoint) ckpt = {k.replace('module.', ''): v for k, v in checkpoint.items()} # strip the names train_model.load_state_dict(ckpt) else: print("=> no checkpoint found at '{}'".format(args.load)) for m in train_model.modules(): if hasattr(m, 'switch_to_deploy'): m.switch_to_deploy() torch.save(train_model.state_dict(), args.save)
def main(arch, model_path, output_path, input_shape=(224, 224), batch_size=1): switch_conv_bn_impl('DBB') switch_deploy_flag(True) model = build_model(arch) model.load_state_dict(torch.load(model_path)) dummy_input = torch.autograd.Variable( torch.randn(batch_size, 3, input_shape[0], input_shape[1])) torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=12) onnx_model = onnx.load(output_path) # load onnx model model_simp, check = simplify(onnx_model) assert check, "Simplified ONNX model could not be validated" onnx.save(model_simp, output_path) print('finished exporting onnx ')
def test(): args = parser.parse_args() switch_deploy_flag(args.mode == 'deploy') switch_conv_bn_impl(args.blocktype) model = build_model(args.arch) if not torch.cuda.is_available(): print('using CPU, this will be slow') use_gpu = False else: model = model.cuda() use_gpu = True # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() if 'hdf5' in args.weights: from utils import model_load_hdf5 model_load_hdf5(model, args.weights) elif os.path.isfile(args.weights): print("=> loading checkpoint '{}'".format(args.weights)) checkpoint = torch.load(args.weights) if 'state_dict' in checkpoint: checkpoint = checkpoint['state_dict'] ckpt = {k.replace('module.', ''): v for k, v in checkpoint.items()} # strip the names model.load_state_dict(ckpt) else: print("=> no checkpoint found at '{}'".format(args.weights)) cudnn.benchmark = True # Data loading code valdir = os.path.join(args.data, 'val') val_loader = torch.utils.data.DataLoader(datasets.ImageFolder( valdir, val_preprocess(224)), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) validate(val_loader, model, criterion, use_gpu)
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) # =========================== build model from convnet_utils import switch_deploy_flag, switch_conv_bn_impl, build_model switch_deploy_flag(False) switch_conv_bn_impl(args.blocktype) model = build_model(args.arch) if gpu == 0: for name, param in model.named_parameters(): print(name, param.size()) 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') trans = strong_train_preprocess(224) if 'ResNet' in args.arch else standard_train_preprocess(224) print('aug is ', trans) train_dataset = datasets.ImageFolder(traindir, trans) 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, val_preprocess(224)) 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, filename='{}_{}.pth.tar'.format(args.arch, args.blocktype))