zoo.examples.orca.learn.horovod package

Submodules

zoo.examples.orca.learn.horovod.pytorch_estimator module

class zoo.examples.orca.learn.horovod.pytorch_estimator.LinearDataset(a, b, size=1000)

Bases: torch.utils.data.dataset.Dataset

y = a * x + b

zoo.examples.orca.learn.horovod.pytorch_estimator.model_creator(config)

Returns a torch.nn.Module object.

zoo.examples.orca.learn.horovod.pytorch_estimator.optimizer_creator(model, config)

Returns optimizer defined upon the model parameters.

zoo.examples.orca.learn.horovod.pytorch_estimator.scheduler_creator(optimizer, config)

Returns a learning rate scheduler wrapping the optimizer. You will need to set TorchTrainer(scheduler_step_freq="epoch") for the scheduler to be incremented correctly. If using a scheduler for validation loss, be sure to call trainer.update_scheduler(validation_loss).

zoo.examples.orca.learn.horovod.pytorch_estimator.train_data_creator(config)

zoo.examples.orca.learn.horovod.pytorch_estimator.train_example(workers_per_node)

zoo.examples.orca.learn.horovod.pytorch_estimator.validation_data_creator(config)

zoo.examples.orca.learn.horovod.simple_horovod_pytorch module

Module contents