From 125f68f305f3a345117b2b95174b7e62eec65654 Mon Sep 17 00:00:00 2001 From: XiaocenDong <63833153+XiaocenDong@users.noreply.github.com> Date: Thu, 23 Apr 2020 23:22:38 +0800 Subject: [PATCH] fixed mnist bug (#3569) * fixed mnist bug * fixed train_step param --- orttraining/orttraining/python/ort_trainer.py | 1 + .../mnist_training.py | 36 ++++++++++++++----- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/orttraining/orttraining/python/ort_trainer.py b/orttraining/orttraining/python/ort_trainer.py index d9f58e0704..c38b92df83 100644 --- a/orttraining/orttraining/python/ort_trainer.py +++ b/orttraining/orttraining/python/ort_trainer.py @@ -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 diff --git a/orttraining/pytorch_frontend_examples/mnist_training.py b/orttraining/pytorch_frontend_examples/mnist_training.py index e7f11e2725..ed73a132e4 100644 --- a/orttraining/pytorch_frontend_examples/mnist_training.py +++ b/orttraining/pytorch_frontend_examples/mnist_training.py @@ -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)