|
| 1 | +# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance with the License. A copy of the License is located at |
| 4 | +# |
| 5 | +# http://aws.amazon.com/apache2.0/ |
| 6 | +# |
| 7 | +# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and limitations under the License. |
| 8 | + |
| 9 | +from __future__ import print_function |
| 10 | +import argparse |
| 11 | +import time |
| 12 | +import torch |
| 13 | +import torch.nn as nn |
| 14 | +import torch.nn.functional as F |
| 15 | +import torch.optim as optim |
| 16 | +from torchvision import datasets, transforms |
| 17 | +from torch.optim.lr_scheduler import StepLR |
| 18 | +from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP |
| 19 | +import smdistributed.dataparallel.torch.distributed as dist |
| 20 | +dist.init_process_group() |
| 21 | + |
| 22 | +class Net(nn.Module): |
| 23 | + def __init__(self): |
| 24 | + super(Net, self).__init__() |
| 25 | + self.conv1 = nn.Conv2d(1, 32, 3, 1) |
| 26 | + self.conv2 = nn.Conv2d(32, 64, 3, 1) |
| 27 | + self.dropout1 = nn.Dropout2d(0.25) |
| 28 | + self.dropout2 = nn.Dropout2d(0.5) |
| 29 | + self.fc1 = nn.Linear(9216, 128) |
| 30 | + self.fc2 = nn.Linear(128, 10) |
| 31 | + |
| 32 | + def forward(self, x): |
| 33 | + x = self.conv1(x) |
| 34 | + x = F.relu(x) |
| 35 | + x = self.conv2(x) |
| 36 | + x = F.relu(x) |
| 37 | + x = F.max_pool2d(x, 2) |
| 38 | + x = self.dropout1(x) |
| 39 | + x = torch.flatten(x, 1) |
| 40 | + x = self.fc1(x) |
| 41 | + x = F.relu(x) |
| 42 | + x = self.dropout2(x) |
| 43 | + x = self.fc2(x) |
| 44 | + output = F.log_softmax(x, dim=1) |
| 45 | + return output |
| 46 | + |
| 47 | + |
| 48 | +def train(args, model, device, train_loader, optimizer, epoch): |
| 49 | + model.train() |
| 50 | + for batch_idx, (data, target) in enumerate(train_loader): |
| 51 | + data, target = data.to(device), target.to(device) |
| 52 | + optimizer.zero_grad() |
| 53 | + output = model(data) |
| 54 | + loss = F.nll_loss(output, target) |
| 55 | + loss.backward() |
| 56 | + optimizer.step() |
| 57 | + if batch_idx % args.log_interval == 0 and args.rank == 0: |
| 58 | + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
| 59 | + epoch, batch_idx * len(data) * args.world_size, len(train_loader.dataset), |
| 60 | + 100. * batch_idx / len(train_loader), loss.item())) |
| 61 | + if args.verbose: |
| 62 | + print('Batch', batch_idx, "from rank", args.rank) |
| 63 | + |
| 64 | + |
| 65 | +def test(model, device, test_loader): |
| 66 | + model.eval() |
| 67 | + test_loss = 0 |
| 68 | + correct = 0 |
| 69 | + with torch.no_grad(): |
| 70 | + for data, target in test_loader: |
| 71 | + data, target = data.to(device), target.to(device) |
| 72 | + output = model(data) |
| 73 | + test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss |
| 74 | + pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability |
| 75 | + correct += pred.eq(target.view_as(pred)).sum().item() |
| 76 | + |
| 77 | + test_loss /= len(test_loader.dataset) |
| 78 | + |
| 79 | + print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( |
| 80 | + test_loss, correct, len(test_loader.dataset), |
| 81 | + 100. * correct / len(test_loader.dataset))) |
| 82 | + |
| 83 | + |
| 84 | +def main(): |
| 85 | + # Training settings |
| 86 | + parser = argparse.ArgumentParser(description='PyTorch MNIST Example') |
| 87 | + parser.add_argument('--batch-size', type=int, default=64, metavar='N', |
| 88 | + help='input batch size for training (default: 64)') |
| 89 | + parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', |
| 90 | + help='input batch size for testing (default: 1000)') |
| 91 | + parser.add_argument('--epochs', type=int, default=14, metavar='N', |
| 92 | + help='number of epochs to train (default: 14)') |
| 93 | + parser.add_argument('--lr', type=float, default=1.0, metavar='LR', |
| 94 | + help='learning rate (default: 1.0)') |
| 95 | + parser.add_argument('--gamma', type=float, default=0.7, metavar='M', |
| 96 | + help='Learning rate step gamma (default: 0.7)') |
| 97 | + parser.add_argument('--seed', type=int, default=1, metavar='S', |
| 98 | + help='random seed (default: 1)') |
| 99 | + parser.add_argument('--log-interval', type=int, default=10, metavar='N', |
| 100 | + help='how many batches to wait before logging training status') |
| 101 | + parser.add_argument('--save-model', action='store_true', default=False, |
| 102 | + help='For Saving the current Model') |
| 103 | + parser.add_argument('--verbose', action='store_true', default=False, |
| 104 | + help='For displaying SMDataParallel-specific logs') |
| 105 | + parser.add_argument('--data-path', type=str, default='/tmp/data', help='Path for downloading ' |
| 106 | + 'the MNIST dataset') |
| 107 | + |
| 108 | + args = parser.parse_args() |
| 109 | + args.world_size = dist.get_world_size() |
| 110 | + args.rank = rank = dist.get_rank() |
| 111 | + args.local_rank = local_rank = dist.get_local_rank() |
| 112 | + args.lr = 1.0 |
| 113 | + args.batch_size //= args.world_size // 8 |
| 114 | + args.batch_size = max(args.batch_size, 1) |
| 115 | + data_path = args.data_path |
| 116 | + |
| 117 | + if args.verbose: |
| 118 | + print('Hello from rank', rank, 'of local_rank', |
| 119 | + local_rank, 'in world size of', args.world_size) |
| 120 | + |
| 121 | + if not torch.cuda.is_available(): |
| 122 | + raise Exception("Must run SMDataParallel MNIST example on CUDA-capable devices.") |
| 123 | + |
| 124 | + torch.manual_seed(args.seed) |
| 125 | + |
| 126 | + device = torch.device("cuda") |
| 127 | + |
| 128 | + if local_rank == 0: |
| 129 | + train_dataset = datasets.MNIST(data_path, train=True, download=True, |
| 130 | + transform=transforms.Compose([ |
| 131 | + transforms.ToTensor(), |
| 132 | + transforms.Normalize((0.1307,), (0.3081,)) |
| 133 | + ])) |
| 134 | + else: |
| 135 | + time.sleep(8) |
| 136 | + train_dataset = datasets.MNIST(data_path, train=True, download=False, |
| 137 | + transform=transforms.Compose([ |
| 138 | + transforms.ToTensor(), |
| 139 | + transforms.Normalize((0.1307,), (0.3081,)) |
| 140 | + ])) |
| 141 | + |
| 142 | + train_sampler = torch.utils.data.distributed.DistributedSampler( |
| 143 | + train_dataset, |
| 144 | + num_replicas=args.world_size, |
| 145 | + rank=rank) |
| 146 | + train_loader = torch.utils.data.DataLoader( |
| 147 | + train_dataset, |
| 148 | + batch_size=args.batch_size, |
| 149 | + shuffle=False, |
| 150 | + num_workers=0, |
| 151 | + pin_memory=True, |
| 152 | + sampler=train_sampler) |
| 153 | + if rank == 0: |
| 154 | + test_loader = torch.utils.data.DataLoader( |
| 155 | + datasets.MNIST(data_path, train=False, transform=transforms.Compose([ |
| 156 | + transforms.ToTensor(), |
| 157 | + transforms.Normalize((0.1307,), (0.3081,)) |
| 158 | + ])), |
| 159 | + batch_size=args.test_batch_size, shuffle=True) |
| 160 | + |
| 161 | + model = DDP(Net().to(device)) |
| 162 | + torch.cuda.set_device(local_rank) |
| 163 | + model.cuda(local_rank) |
| 164 | + optimizer = optim.Adadelta(model.parameters(), lr=args.lr) |
| 165 | + scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) |
| 166 | + for epoch in range(1, args.epochs + 1): |
| 167 | + train(args, model, device, train_loader, optimizer, epoch) |
| 168 | + if rank == 0: |
| 169 | + test(model, device, test_loader) |
| 170 | + scheduler.step() |
| 171 | + |
| 172 | + if args.save_model: |
| 173 | + torch.save(model.state_dict(), "mnist_cnn.pt") |
| 174 | + |
| 175 | + |
| 176 | +if __name__ == '__main__': |
| 177 | + main() |
0 commit comments