def train_mnist():
model = MLP()
loss_func = F.nll_loss
device = "cuda:0"
batch_size = 5
trainer = Trainer(model, loss_func, device=device)
train_loader, _ = data_loader.get_mnist_dataloader(batch_size=batch_size,
shuffle=False)
print("feeding")
i = 0
for i, batch in enumerate(train_loader):
feed_data = [t.to(device) for t in batch]
pytorch_loss = trainer.run_by_pytorch(*feed_data)
nnf_loss = trainer(*feed_data)
if i % 100 == 0:
print("iter ", i)
print('pytorch_loss: ', pytorch_loss)
print('nnf_loss: ', nnf_loss)
if i == 10000:
break
torch.save(model.state_dict(), "/tmp/mnist.pt")
def eval_rnn():
print("Evaluate RNN model on MNIST dataset")
train_loader, test_loader = data_loader.get_mnist_dataloader(
batch_size=args.batch_size,
shuffle=False,
drop_last=True,
num_workers=2)
model = ImageClassification(28, 10, args.num_layer, SimpleRNNCell,
args.hidden_size, args.batch_size, 28)
state_dict = torch.load("/tmp/rnn.pt")
model.load_state_dict(state_dict)
model.to(args.device)
model.eval()
eval_acc = 0
eval_size = 0
for i, batch in enumerate(test_loader):
with torch.no_grad():
data, label = (t.to(args.device) for t in batch)
predict = model(data)
eval_acc += (predict.max(dim=1)[1] == label).sum().cpu().numpy()
eval_size += args.batch_size
print("Evaluation accuracy:", (eval_acc / eval_size) * 100, "%")
def pytorch_train_rnn():
backend = "PyTorch"
if args.backend == "torchscript":
backend += " TorchScript"
print("Train RNN model on MNIST dataset with {}".format(backend))
train_loader, test_loader = data_loader.get_mnist_dataloader(
batch_size=args.batch_size,
shuffle=False,
drop_last=True,
num_workers=2)
model = ImageClassification(28, 10, args.num_layer, SimpleRNNCell,
args.hidden_size, args.batch_size,
28).to(args.device)
if args.backend == "torchscript":
model = torch.jit.trace(
model, (torch.ones(args.batch_size, 1, 28, 28).to(args.device)))
optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)
loss_func = F.nll_loss
sum_loss = 0
sum_iter = 0
sum_time = 0
torch.cuda.synchronize()
start_time = time.time()
for epoch in range(args.num_epoch):
for i, batch in enumerate(train_loader):
if sum_iter == 100:
print("Epoch {}, batch {}，loss {}, time {} s".format(
epoch, i, sum_loss / sum_iter, sum_time / sum_iter))
sum_loss = 0
sum_iter = 0
sum_time = 0
data, label = (t.to(args.device) for t in batch)
predict = model(data)
loss = loss_func(predict, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
torch.cuda.synchronize()
end_time = time.time()
iter_time = end_time - start_time
start_time = end_time
sum_loss += loss
sum_time += iter_time
sum_iter += 1
torch.save(model.state_dict(), "/tmp/rnn.pt")
def eval():
device = "cuda:0"
model = MLP()
state_dict = torch.load("/tmp/mnist.pt")
model.load_state_dict(state_dict)
model.to(device)
model.eval()
_, test_loader = data_loader.get_mnist_dataloader(batch_size=1,
shuffle=False)
for i, batch in enumerate(test_loader):
with torch.no_grad():
pred = model(batch[0].to(device))
prob = np_softmax(pred.cpu().numpy())
print('Image {} is digit "{}", confidence {}'.format(
i, np.argmax(prob), np.max(prob)))
if i == 5:
break
def train_mnist():
model = MLP()
loss_func = F.nll_loss
device = "cuda:0"
batch_size = 5
codegen_flags = {
"autodiff":
True, # add backward graph
"training_mode":
True, # move weight external
"extern_result_memory":
True, # move result external
"training_optimizer":
'\'' + json.dumps({
"optimizer": "SGD",
"learning_rate": 0.001
}) + '\'', # training optimizer configs
}
trainer = Trainer(model,
loss_func,
device=device,
codegen_flags=codegen_flags)
train_loader, _ = data_loader.get_mnist_dataloader(batch_size=batch_size,
shuffle=False)
print("feeding")
i = 0
for i, batch in enumerate(train_loader):
feed_data = [t.to(device) for t in batch]
pytorch_loss = trainer.run_by_pytorch(*feed_data)
nnf_loss = trainer(*feed_data)
if i % 100 == 0:
print("iter ", i)
print('pytorch_loss: ', pytorch_loss)
print('nnf_loss: ', nnf_loss)
if i == 10000:
break
torch.save(model.state_dict(), "/tmp/mnist.pt")
def train_rnn():
print("Train RNN model on MNIST dataset with NNFusion")
train_loader, test_loader = data_loader.get_mnist_dataloader(
batch_size=args.batch_size,
shuffle=False,
drop_last=True,
num_workers=2)
model = ImageClassification(28, 10, args.num_layer, SimpleRNNCell,
args.hidden_size, args.batch_size,
28).to(args.device)
loss_func = F.nll_loss
codegen_flags = {
"autodiff":
True, # add backward graph
"training_mode":
True, # move weight external
"extern_result_memory":
True, # move result external
"kernels_as_files":
True, # enable parallel compilation for generated model code
"kernels_files_number":
20, # number of parallel compilation
"training_optimizer":
'\'' + json.dumps({
"optimizer": "SGD",
"learning_rate": args.learning_rate
}) + '\'', # training optimizer configs
}
trainer = Trainer(model,
loss_func,
device=args.device,
codegen_flags=codegen_flags)
sum_loss = 0
sum_iter = 0
sum_time = 0
torch.cuda.synchronize()
start_time = time.time()
for epoch in range(args.num_epoch):
for i, batch in enumerate(train_loader):
if sum_iter == 100:
print("Epoch {}, batch {}，loss {}, time {} s".format(
epoch, i, sum_loss / sum_iter, sum_time / sum_iter))
sum_loss = 0
sum_iter = 0
sum_time = 0
data, label = (t.to(args.device) for t in batch)
nnf_loss = trainer(data, label)
torch.cuda.synchronize()
end_time = time.time()
iter_time = end_time - start_time
start_time = end_time
sum_loss += nnf_loss
sum_time += iter_time
sum_iter += 1
torch.save(model.state_dict(), "/tmp/rnn.pt")