def setup_net(snapshot): """Quickly create a network for the given snapshot. Arguments: snapshot {string} -- Input snapshot, IE. kitti_best.pth Returns: [net, transform] -- PyTorch model & the image transform function. """ cudnn.benchmark = False torch.cuda.empty_cache() args = Args('./save', 'network.deepv3.DeepWV3Plus', snapshot) assert_and_infer_cfg(args, train_mode=False) # get net net = network.get_net(args, criterion=None) net = torch.nn.DataParallel(net).cuda() print('Net built.') net, _ = restore_snapshot(net, optimizer=None, snapshot=snapshot, restore_optimizer_bool=False) net.eval() print('Net restored.') # get data mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) img_transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize(*mean_std)]) return net, img_transform, args
def citymode(): model = deepv3.DeepWV3Plus(21) model_path = rep_path + '/pretrained_models/cityscapes_best.pth' model, _ = restore_snapshot(model, optimizer=None, snapshot=model_path, restore_optimizer_bool=False) model.cpu() return model
def get_net(): """ Get Network for evaluation """ logging.info('Load model file: %s', args.snapshot) net = network.get_net_ori(args, criterion=None) net = torch.nn.DataParallel(net).cuda() net, _ = restore_snapshot(net, optimizer=None, snapshot=args.snapshot, snapshot2=args.snapshot2, restore_optimizer_bool=False) net.eval() return net
def get_net(optimizer=None, criterion=None): """ Get network for train :return: """ net = network.get_net(args, criterion=criterion) net, _ = restore_snapshot(net, optimizer=optimizer, snapshot=args.snapshot, restore_optimizer_bool=False) #net.train() return net
def get_net(): """ Get Network for evaluation """ logging.info('Load model file: %s', args.snapshot) net = network.get_net(args, criterion=None) net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net) net = network.warp_network_in_dataparallel(net, args.local_rank) net, _, _, _, _ = restore_snapshot(net, optimizer=None, scheduler=None, snapshot=args.snapshot, restore_optimizer_bool=False) net.eval() return net
def get_net(): """ Get Network for evaluation """ logging.info('Load model file: %s', args.snapshot) net = network.get_net(args, criterion=None) if args.inference_mode == 'pooling': net = MyDataParallel(net, gather=False).cuda() else: net = torch.nn.DataParallel(net).cuda() net, _ = restore_snapshot(net, optimizer=None, snapshot=args.snapshot, restore_optimizer_bool=False) net.eval() return net
def get_segmentation(self): # Get Segmentation Net assert_and_infer_cfg(self.opt, train_mode=False) self.opt.dataset_cls = cityscapes net = network.get_net(self.opt, criterion=None) net = torch.nn.DataParallel(net).cuda() print('Segmentation Net Built.') snapshot = os.path.join(os.getcwd(), os.path.dirname(__file__), self.opt.snapshot) self.seg_net, _ = restore_snapshot(net, optimizer=None, snapshot=snapshot, restore_optimizer_bool=False) self.seg_net.eval() print('Segmentation Net Restored.')
parser = argparse.ArgumentParser(description='demo') parser.add_argument('--demo-image', type=str, default='', help='path to demo image', required=True) parser.add_argument('--snapshot', type=str, default='./pretrained_models/cityscapes_best_wideresnet38.pth', help='pre-trained checkpoint', required=True) parser.add_argument('--arch', type=str, default='network.deepv3.DeepWV3Plus', help='network architecture used for inference') parser.add_argument('--save-dir', type=str, default='./save', help='path to save your results') args = parser.parse_args() assert_and_infer_cfg(args, train_mode=False) cudnn.benchmark = False torch.cuda.empty_cache() # get net args.dataset_cls = cityscapes net = network.get_net(args, criterion=None) net = torch.nn.DataParallel(net).cuda() print('Net built.') net, _ = restore_snapshot(net, optimizer=None, snapshot=args.snapshot, restore_optimizer_bool=False) net.eval() print('Net restored.') # get data mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) img_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(*mean_std)]) img = Image.open(args.demo_image).convert('RGB') img_tensor = img_transform(img) # predict with torch.no_grad(): img = img_tensor.unsqueeze(0).cuda() pred = net(img) print('Inference done.')
def convert_segmentation_model(model_name='segmentation.onnx'): assert_and_infer_cfg(opt, train_mode=False) cudnn.benchmark = False torch.cuda.empty_cache() # Get segmentation Net opt.dataset_cls = cityscapes net = network.get_net(opt, criterion=None) net = torch.nn.DataParallel(net).cuda() print('Segmentation Net built.') net, _ = restore_snapshot(net, optimizer=None, snapshot=opt.snapshot, restore_optimizer_bool=False) net.eval() print('Segmentation Net Restored.') # Input to the model batch_size = 1 x = torch.randn(batch_size, 3, 1024, 2048, requires_grad=True).cuda() torch_out = net(x) # Export the model torch.onnx.export( net.module, # model being run x, # model input (or a tuple for multiple inputs) model_name, # where to save the model (can be a file or file-like object) export_params= True, # store the trained parameter weights inside the model file opset_version=11, # the ONNX version to export the model to do_constant_folding= True, # whether to execute constant folding for optimization input_names=['input'], # the model's input names output_names=['output'], # the model's output names dynamic_axes={ 'input': { 0: 'batch_size' }, # variable lenght axes 'output': { 0: 'batch_size' } }) ort_session = onnxruntime.InferenceSession(model_name) def to_numpy(tensor): return tensor.detach().cpu().numpy( ) if tensor.requires_grad else tensor.cpu().numpy() ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)} ort_outs = ort_session.run(None, ort_inputs) # compare ONNX Runtime and PyTorch results np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-03) print( "Exported model has been tested with ONNXRuntime, and the result looks good!" )