def model_fn(model_dir):
print("In model_fn. Model directory is -")
print(model_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net()
with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
print("Loading the mnist model")
model.load_state_dict(torch.load(f, map_location=device))
return model
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size',
type=int,
default=1000,
metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs',
type=int,
default=14,
metavar='N',
help='number of epochs to train (default: 14)')
parser.add_argument('--lr',
type=float,
default=1.0,
metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma',
type=float,
default=0.7,
metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument(
'--log-interval',
type=int,
default=10,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model',
action='store_true',
default=False,
help='For Saving the current Model')
parser.add_argument('--verbose',
action='store_true',
default=False,
help='For displaying SMDataParallel-specific logs')
parser.add_argument('--data-path',
type=str,
default='/tmp/data',
help='Path for downloading '
'the MNIST dataset')
# Model checkpoint location
parser.add_argument('--model-dir',
type=str,
default=os.environ['SM_MODEL_DIR'])
args = parser.parse_args()
args.world_size = dist.get_world_size()
args.rank = rank = dist.get_rank()
args.local_rank = local_rank = dist.get_local_rank()
args.lr = 1.0
args.batch_size //= args.world_size // 8
args.batch_size = max(args.batch_size, 1)
data_path = args.data_path
if args.verbose:
print('Hello from rank', rank, 'of local_rank', local_rank,
'in world size of', args.world_size)
if not torch.cuda.is_available():
raise Exception(
"Must run SMDataParallel MNIST example on CUDA-capable devices.")
torch.manual_seed(args.seed)
device = torch.device("cuda")
if local_rank == 0:
train_dataset = datasets.MNIST(data_path,
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ),
(0.3081, ))
]))
else:
# TODO: Reduce time to half when upgrade to torchvision==0.9.1
time.sleep(16)
train_dataset = datasets.MNIST(data_path,
train=True,
download=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ),
(0.3081, ))
]))
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset, num_replicas=args.world_size, rank=rank)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=0,
pin_memory=True,
sampler=train_sampler)
if rank == 0:
test_loader = torch.utils.data.DataLoader(
datasets.MNIST(data_path,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=args.test_batch_size,
shuffle=True)
# Use SMDataParallel PyTorch DDP for efficient distributed training
model = DDP(Net().to(device))
torch.cuda.set_device(local_rank)
model.cuda(local_rank)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
if rank == 0:
test(model, device, test_loader)
scheduler.step()
if rank == 0:
print("Saving the model...")
save_model(model, args.model_dir)
def main():
# Training settings
parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
parser.add_argument(
"--batch-size",
type=int,
default=64,
metavar="N",
help="input batch size for training (default: 64)",
)
parser.add_argument(
"--test-batch-size",
type=int,
default=1000,
metavar="N",
help="input batch size for testing (default: 1000)",
)
parser.add_argument(
"--epochs",
type=int,
default=14,
metavar="N",
help="number of epochs to train (default: 14)",
)
parser.add_argument("--lr",
type=float,
default=1.0,
metavar="LR",
help="learning rate (default: 1.0)")
parser.add_argument(
"--gamma",
type=float,
default=0.7,
metavar="M",
help="Learning rate step gamma (default: 0.7)",
)
parser.add_argument("--seed",
type=int,
default=1,
metavar="S",
help="random seed (default: 1)")
parser.add_argument(
"--log-interval",
type=int,
default=10,
metavar="N",
help="how many batches to wait before logging training status",
)
parser.add_argument("--save-model",
action="store_true",
default=False,
help="For Saving the current Model")
parser.add_argument(
"--verbose",
action="store_true",
default=False,
help="For displaying smdistributed.dataparallel-specific logs",
)
parser.add_argument(
"--data-path",
type=str,
default="/tmp/data",
help="Path for downloading "
"the MNIST dataset",
)
args = parser.parse_args()
args.world_size = dist.get_world_size()
args.rank = rank = dist.get_rank()
args.local_rank = local_rank = dist.get_local_rank()
args.lr = 1.0
args.batch_size //= args.world_size // 8
args.batch_size = max(args.batch_size, 1)
data_path = args.data_path
if args.verbose:
print(
"Hello from rank",
rank,
"of local_rank",
local_rank,
"in world size of",
args.world_size,
)
if not torch.cuda.is_available():
raise CUDANotFoundException(
"Must run smdistributed.dataparallel MNIST example on CUDA-capable devices."
)
torch.manual_seed(args.seed)
device = torch.device("cuda")
# select a single rank per node to download data
is_first_local_rank = local_rank == 0
if is_first_local_rank:
train_dataset = datasets.MNIST(
data_path,
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
]),
)
dist.barrier() # prevent other ranks from accessing the data early
if not is_first_local_rank:
train_dataset = datasets.MNIST(
data_path,
train=True,
download=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
]),
)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset, num_replicas=args.world_size, rank=rank)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=0,
pin_memory=True,
sampler=train_sampler,
)
if rank == 0:
test_loader = torch.utils.data.DataLoader(
datasets.MNIST(
data_path,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
]),
),
batch_size=args.test_batch_size,
shuffle=True,
)
model = DDP(Net().to(device))
torch.cuda.set_device(local_rank)
model.cuda(local_rank)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
if rank == 0:
test(model, device, test_loader)
scheduler.step()
if args.save_model:
torch.save(model.state_dict(), "mnist_cnn.pt")
def train(args):
logger.debug("Number of gpus available - {}".format(args.num_gpus))
# Horovod: initialize library
hvd.init()
torch.manual_seed(args.seed)
# Horovod: pin GPU to local local rank
torch.cuda.set_device(hvd.local_rank())
torch.cuda.manual_seed(args.seed)
# Horovod: limit number of CPU threads to be used per worker
torch.set_num_threads(1)
kwargs = {'num_workers': 1, 'pin_memory': True}
train_loader = _get_train_data_loader(args.batch_size, args.data_dir,
**kwargs)
test_loader = _get_test_data_loader(args.test_batch_size, args.data_dir,
**kwargs)
logger.debug("Processes {}/{} ({:.0f}%) of train data".format(
len(train_loader.sampler), len(train_loader.dataset),
100. * len(train_loader.sampler) / len(train_loader.dataset)))
logger.debug("Processes {}/{} ({:.0f}%) of test data".format(
len(test_loader.sampler), len(test_loader.dataset),
100. * len(test_loader.sampler) / len(test_loader.dataset)))
model = Net()
lr_scaler = hvd.size()
model.cuda()
# Horovod: scale learning rate by lr_scaler.
optimizer = optim.SGD(model.parameters(),
lr=args.lr * lr_scaler,
momentum=args.momentum)
# Horovod: broadcast parameters & optimizer state.
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
# Horovod: wrap optimizer with DistributedOptimizer.
optimizer = hvd.DistributedOptimizer(
optimizer, named_parameters=model.named_parameters())
for epoch in range(1, args.epochs + 1):
model.train()
for batch_idx, (data, target) in enumerate(train_loader, 1):
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
logger.info(
'Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(
epoch, batch_idx * len(data),
len(train_loader.sampler),
100. * batch_idx / len(train_loader), loss.item()))
test(model, test_loader)
save_model(model, args.model_dir)
def model_fn(model_dir):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net()
with open(os.path.join(model_dir, "model.pth"), "rb") as f:
model.load_state_dict(torch.load(f))
return model.to(device)