def forward_pytorch(weightfile, image): #net=resnet.resnet18() #net = resnet.resnet18() net=LeNet(1,2) checkpoint = torch.load(weightfile) net.load_state_dict(checkpoint['weight']) net.double() # to double if args.cuda: net.cuda() print(net) net.eval() image = torch.from_numpy(image.astype(np.float64)) # to double if args.cuda: image = Variable(image.cuda()) else: image = Variable(image) t0 = time.time() blobs = net.forward(image) print(blobs.data.numpy().flatten()) t1 = time.time() return t1-t0, blobs, net, torch.from_numpy(blobs.data.numpy())
def main(): global args, rank, world_size, best_prec1, dataset_len if args.dist == 1: rank, world_size = dist_init() else: rank = 0 world_size = 1 model = LeNet() model.cuda() if args.double == 1: param_copy = [ param.clone().type(torch.cuda.DoubleTensor).detach() for param in model.parameters() ] else: param_copy = [ param.clone().type(torch.cuda.FloatTensor).detach() for param in model.parameters() ] for param in param_copy: param.requires_grad = True if args.double == 1: model = model.double() if args.half == 1: model = model.half() if args.dist == 1: model = DistModule(model) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(param_copy, args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint last_iter = -1 # Data loading code train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=False) val_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=False) dataset_len = len(train_dataset) args.max_iter = math.ceil( (dataset_len * args.epoch) / (world_size * args.batch_size)) if args.dist == 1: train_sampler = DistributedGivenIterationSampler(train_dataset, args.max_iter, args.batch_size, last_iter=last_iter) val_sampler = DistributedSampler(val_dataset, round_up=False) else: train_sampler = DistributedGivenIterationSampler(train_dataset, args.max_iter, args.batch_size, world_size=1, rank=0, last_iter=last_iter) val_sampler = None # pin_memory if true, will copy the tensor to cuda pinned memory train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler) train(train_loader, val_loader, model, criterion, optimizer, param_copy)