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)
