saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fixed mnist bug (#3569)

Browse source 
* fixed mnist bug

* fixed train_step param
This commit is contained in:
XiaocenDong 2020-04-23 23:22:38 +08:00 committed by GitHub
parent f1ba9aaf34
commit 125f68f305
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
orttraining/orttraining/python/ort_trainer.py
View file

@ -510,6 +510,7 @@ def create_ort_training_session_bind_parameters(model, device, world_rank=-1, wo
dtype_torch_to_numpy(torch_params[param].dtype), list(torch_tensor.size()),
torch_tensor.data_ptr())
device_index = get_device_index(device)
create_and_bind_grad_or_grad_accumulate_buffer(train_io_binding, torch_tensor, param, enable_grad_accumulation, device, device_index)
return session, train_io_binding, eval_io_binding, output_name, torch_params, output_types

36
orttraining/pytorch_frontend_examples/mnist_training.py
View file

@ -17,6 +17,7 @@ import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import os
# TODO: remove after ready for CV
# import sys
@ -27,13 +28,15 @@ import numpy as np
# from ort_trainer import IODescription, ModelDescription, ORTTrainer, ORTModel
from onnxruntime.capi.ort_trainer import IODescription, ModelDescription, ORTTrainer, ORTModel
from mpi4py import MPI
from onnxruntime.capi._pybind_state import set_cuda_device_id
class NeuralNet(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(NeuralNet, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
@ -79,18 +82,19 @@ def test_with_model(args, model, device, test_loader, optimizer, epoch):
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
def train_with_trainer(args, trainer, device, train_loader, epoch):
def train_with_trainer(args, trainer, device, train_loader, epoch):
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
data = data.reshape(data.shape[0], -1)
learning_rate = torch.tensor([args.lr])
loss = trainer.train_step((data, target, learning_rate))
loss = trainer.train_step(data, target, learning_rate)
# Since the output corresponds to [loss_desc, probability_desc], the first value is taken as loss.
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
100. * batch_idx / len(train_loader), loss[0]))
# TODO: comple this once ORT training can do evaluation.
def test_with_trainer(args, trainer, device, test_loader):
@ -152,8 +156,6 @@ def main():
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 0, 'pin_memory': True}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=True, download=True,
@ -167,6 +169,19 @@ def main():
transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])),
batch_size=args.test_batch_size, shuffle=True, **kwargs)
comm = MPI.COMM_WORLD
args.local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) if ('OMPI_COMM_WORLD_LOCAL_RANK' in os.environ) else 0
args.world_rank = int(os.environ['OMPI_COMM_WORLD_RANK']) if ('OMPI_COMM_WORLD_RANK' in os.environ) else 0
args.world_size=comm.Get_size()
torch.cuda.set_device(args.local_rank)
if use_cuda:
device = torch.device("cuda", args.local_rank)
else:
device = torch.device("cpu")
args.n_gpu = 1
set_cuda_device_id(args.local_rank)
input_size = 784
hidden_size = 500
num_classes = 10
@ -175,14 +190,17 @@ def main():
model_desc = mnist_model_description()
if args.use_ort_trainer:
# use log_interval as gradient accumulate steps
trainer = ORTTrainer(model, my_loss, model_desc, "SGDOptimizer", None, IODescription('Learning_Rate', [1,], torch.float32), device)
trainer = ORTTrainer(model, my_loss, model_desc, "LambOptimizer", None, IODescription('Learning_Rate', [1,], torch.float32), device, 1, None,
args.world_rank, args.world_size, use_mixed_precision=False, allreduce_post_accumulation = True)
print('\nBuild ort model done.')
for epoch in range(1, args.epochs + 1):
train_with_trainer(args, trainer, device, train_loader, epoch)
import pdb
test_with_trainer(args, trainer, device, test_loader)
else:
model = ORTModel(model, my_loss, model_desc, device)
model = ORTModel(model, my_loss, model_desc, device, None, args.world_rank, args.world_size)
print('\nBuild ort model done.')
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)