def main(run_id, dataset_dir, ckpt_run_dir, output_dir, validation_only=False):
r"""Main logic."""
num_parallel_workers = 0
use_cuda = True
max_batch_per_epoch = None
train_epochs = 164
batch_size = 128
initialize_backends(comm_backend='mpi',
logging_level='INFO',
logging_file=os.path.join(output_dir, 'mlbench.log'),
use_cuda=use_cuda,
seed=42,
cudnn_deterministic=False,
ckpt_run_dir=ckpt_run_dir,
delete_existing_ckpts=not validation_only)
rank = dist.get_rank()
world_size = dist.get_world_size()
model = ResNetCIFAR(resnet_size=20,
bottleneck=False,
num_classes=10,
version=1)
ring_neighbors = [(rank + 1) % world_size, (rank - 1) % world_size]
optimizer = DecentralizedSGD(rank=rank,
neighbors=ring_neighbors,
model=model,
lr=0.1,
momentum=0.9,
weight_decay=1e-4,
nesterov=False)
# Create a learning rate scheduler for an optimizer
scheduler = MultiStepLR(optimizer, milestones=[82, 109], gamma=0.1)
# A loss_function for computing the loss
loss_function = CrossEntropyLoss()
if use_cuda:
model = model.cuda()
loss_function = loss_function.cuda()
# Metrics like Top 1/5 Accuracy
metrics = [TopKAccuracy(topk=1), TopKAccuracy(topk=5)]
train_set = CIFAR10V1(dataset_dir, train=True, download=True)
val_set = CIFAR10V1(dataset_dir, train=False, download=True)
train_set = partition_dataset_by_rank(train_set, rank, world_size)
val_set = partition_dataset_by_rank(val_set, rank, world_size)
train_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=num_parallel_workers,
pin_memory=use_cuda,
drop_last=False)
val_loader = DataLoader(val_set,
batch_size=batch_size,
shuffle=False,
num_workers=num_parallel_workers,
pin_memory=use_cuda,
drop_last=False)
checkpointer = Checkpointer(ckpt_run_dir=ckpt_run_dir,
rank=rank,
freq=CheckpointFreq.NONE)
if not validation_only:
controlflow = TrainValidation(model=model,
optimizer=optimizer,
loss_function=loss_function,
metrics=metrics,
scheduler=scheduler,
batch_size=batch_size,
train_epochs=train_epochs,
rank=rank,
world_size=world_size,
run_id=run_id,
dtype='fp32',
validate=True,
schedule_per='epoch',
checkpoint=checkpointer,
transform_target_type=None,
average_models=True,
use_cuda=use_cuda,
max_batch_per_epoch=max_batch_per_epoch)
controlflow.run(dataloader_train=train_loader,
dataloader_val=val_loader,
dataloader_train_fn=None,
dataloader_val_fn=None,
resume=False,
repartition_per_epoch=False)
else:
cecf = CheckpointsEvaluationControlFlow(ckpt_dir=ckpt_run_dir,
rank=rank,
world_size=world_size,
checkpointer=checkpointer,
model=model,
epochs=train_epochs,
loss_function=loss_function,
metrics=metrics,
use_cuda=use_cuda,
dtype='fp32',
max_batch_per_epoch=None)
train_stats = cecf.evaluate_by_epochs(train_loader)
with open(os.path.join(output_dir, "train_stats.json"), 'w') as f:
json.dump(train_stats, f)
val_stats = cecf.evaluate_by_epochs(val_loader)
with open(os.path.join(output_dir, "val_stats.json"), 'w') as f:
json.dump(val_stats, f)
def main(run_id):
checkpoint_dir = os.path.join(config['checkpoint_root'], run_id)
rank, world_size, _ = initialize_backends(
comm_backend=config['comm_backend'],
logging_level=config['logging_level'],
logging_file=config['logging_file'],
use_cuda=config['use_cuda'],
seed=config['seed'],
ckpt_run_dir=checkpoint_dir)
os.makedirs(config['dataset_root'], exist_ok=True)
train_set = CIFAR10V1(config['dataset_root'], train=True, download=True)
val_set = CIFAR10V1(config['dataset_root'], train=False, download=True)
train_set = partition_dataset_by_rank(train_set, rank, world_size)
train_loader = DataLoader(train_set,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_parallel_workers'],
pin_memory=config['use_cuda'],
drop_last=False)
val_loader = DataLoader(val_set,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_parallel_workers'],
pin_memory=config['use_cuda'],
drop_last=False)
model = get_resnet_model('resnet20',
2,
'fp32',
num_classes=config['num_classes'],
use_cuda=True)
if config['use_cuda']:
model.cuda()
lr = config['lr_per_sample'] * config['batch_size']
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=config['momentum'],
weight_decay=config['weight_decay'],
nesterov=config['nesterov'])
scheduler = multistep_learning_rates_with_warmup(
optimizer,
world_size,
lr,
config['multisteplr_gamma'],
config['multisteplr_milestones'],
warmup_duration=config['warmup_duration'],
warmup_linear_scaling=config['warmup_linear_scaling'],
warmup_lr=lr)
loss_function = CrossEntropyLoss()
if config['use_cuda']:
loss_function.cuda()
metrics = [TopKAccuracy(topk=1), TopKAccuracy(topk=5)]
checkpointer = Checkpointer(checkpoint_dir, rank)
controlflow = TrainValidation(model,
optimizer,
loss_function,
metrics,
scheduler,
config['batch_size'],
config['train_epochs'],
rank,
world_size,
run_id,
dtype=config['dtype'],
checkpoint=checkpointer,
use_cuda=config['use_cuda'])
controlflow.run(dataloader_train=train_loader, dataloader_val=val_loader)
def train_loop(run_id,
dataset_dir,
ckpt_run_dir,
output_dir,
validation_only=False,
use_cuda=False,
light_target=False):
r"""Main logic."""
num_parallel_workers = 2
max_batch_per_epoch = None
train_epochs = 164
batch_size = 128
rank = dist.get_rank()
world_size = dist.get_world_size()
model = ResNetCIFAR(resnet_size=20,
bottleneck=False,
num_classes=10,
version=1)
optimizer = CentralizedSGD(world_size=world_size,
model=model,
lr=0.1,
momentum=0.9,
weight_decay=1e-4,
nesterov=False)
# Create a learning rate scheduler for an optimizer
scheduler = MultiStepLR(optimizer, milestones=[82, 109], gamma=0.1)
# A loss_function for computing the loss
loss_function = CrossEntropyLoss()
if use_cuda:
model = model.cuda()
loss_function = loss_function.cuda()
# Metrics like Top 1/5 Accuracy
metrics = [TopKAccuracy(topk=1), TopKAccuracy(topk=5)]
train_set = CIFAR10V1(dataset_dir, train=True, download=True)
val_set = CIFAR10V1(dataset_dir, train=False, download=True)
train_set = partition_dataset_by_rank(train_set, rank, world_size)
val_set = partition_dataset_by_rank(val_set, rank, world_size)
train_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=num_parallel_workers,
pin_memory=use_cuda,
drop_last=False)
val_loader = DataLoader(val_set,
batch_size=batch_size,
shuffle=False,
num_workers=num_parallel_workers,
pin_memory=use_cuda,
drop_last=False)
checkpointer = Checkpointer(ckpt_run_dir=ckpt_run_dir,
rank=rank,
freq=CheckpointFreq.NONE)
if not validation_only:
if light_target:
goal = task1_time_to_accuracy_light_goal
else:
goal = task1_time_to_accuracy_goal
tracker = Tracker(metrics, run_id, rank, goal=goal)
dist.barrier()
tracker.start()
for epoch in range(0, train_epochs):
train_round(train_loader,
model,
optimizer,
loss_function,
metrics,
scheduler,
'fp32',
schedule_per='epoch',
transform_target_type=None,
use_cuda=use_cuda,
max_batch_per_epoch=max_batch_per_epoch,
tracker=tracker)
is_best = validation_round(val_loader,
model,
loss_function,
metrics,
run_id,
rank,
'fp32',
transform_target_type=None,
use_cuda=use_cuda,
max_batch_per_epoch=max_batch_per_epoch,
tracker=tracker)
checkpointer.save(tracker, model, optimizer, scheduler,
tracker.current_epoch, is_best)
tracker.epoch_end()
if tracker.goal_reached:
print("Goal Reached!")
return
else:
cecf = CheckpointsEvaluationControlFlow(ckpt_dir=ckpt_run_dir,
rank=rank,
world_size=world_size,
checkpointer=checkpointer,
model=model,
epochs=train_epochs,
loss_function=loss_function,
metrics=metrics,
use_cuda=use_cuda,
dtype='fp32',
max_batch_per_epoch=None)
train_stats = cecf.evaluate_by_epochs(train_loader)
with open(os.path.join(output_dir, "train_stats.json"), 'w') as f:
json.dump(train_stats, f)
val_stats = cecf.evaluate_by_epochs(val_loader)
with open(os.path.join(output_dir, "val_stats.json"), 'w') as f:
json.dump(val_stats, f)
def main(run_id):
r"""Main logic."""
num_parallel_workers = 2
dataset_root = '/datasets/torch/cifar10'
use_cuda = True
batch_size = 128
initialize_backends(comm_backend='mpi',
logging_level='INFO',
logging_file='/mlbench.log',
use_cuda=use_cuda,
seed=42,
cudnn_deterministic=False,
ckpt_run_dir='/checkpoints',
delete_existing_ckpts=False)
rank = dist.get_rank()
world_size = dist.get_world_size()
model = ResNetCIFAR(resnet_size=20,
bottleneck=False,
num_classes=10,
version=1)
optimizer = optim.SGD(model.parameters(),
lr=0.1,
momentum=0.9,
weight_decay=1e-4,
nesterov=True)
# Create a learning rate scheduler for an optimizer
scheduler = MultiStepLR(optimizer,
milestones=[82, 109],
gamma=0.1)
# A loss_function for computing the loss
loss_function = CrossEntropyLoss()
if use_cuda:
model = model.cuda()
loss_function = loss_function.cuda()
# Metrics like Top 1/5 Accuracy
metrics = [TopKAccuracy(topk=1), TopKAccuracy(topk=5)]
train_set = CIFAR10V1(dataset_root, train=True, download=True)
val_set = CIFAR10V1(dataset_root, train=False, download=True)
train_set = partition_dataset_by_rank(train_set, rank, world_size)
train_loader = DataLoader(
train_set, batch_size=batch_size, shuffle=True,
num_workers=num_parallel_workers,
pin_memory=use_cuda, drop_last=False)
val_loader = DataLoader(
val_set, batch_size=batch_size, shuffle=False,
num_workers=num_parallel_workers,
pin_memory=use_cuda, drop_last=False)
checkpointer = Checkpointer(
ckpt_run_dir='/checkpoints',
rank=rank,
checkpoint_all=True)
controlflow = TrainValidation(
model=model,
optimizer=optimizer,
loss_function=loss_function,
metrics=metrics,
scheduler=scheduler,
batch_size=batch_size,
train_epochs=164,
rank=rank,
world_size=world_size,
run_id=run_id,
dtype='fp32',
validate=True,
schedule_per='epoch',
checkpoint=checkpointer,
transform_target_type=None,
average_models=True,
use_cuda=True,
max_batch_per_epoch=None)
controlflow.run(
dataloader_train=train_loader,
dataloader_val=val_loader,
dataloader_train_fn=None,
dataloader_val_fn=None,
resume=False,
repartition_per_epoch=False)