def train_mnist(config, checkpoint_dir=False): use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") train_loader, test_loader = get_data_loaders() model = ConvNet().to(device) optimizer = optim.SGD(model.parameters(), lr=0.1) if checkpoint_dir: with open(os.path.join(checkpoint_dir, "checkpoint")) as f: model_state, optimizer_state = torch.load(f) model.load_state_dict(model_state) optimizer.load_state_dict(optimizer_state) model = DistributedDataParallel(model) for epoch in range(40): train(model, optimizer, train_loader, device) acc = test(model, test_loader, device) if epoch % 3 == 0: with distributed_checkpoint_dir(step=epoch) as checkpoint_dir: path = os.path.join(checkpoint_dir, "checkpoint") torch.save((model.state_dict(), optimizer.state_dict()), path) tune.report(mean_accuracy=acc)
def test_checkpoint(ray_start_2_cpus, rank): # noqa: F811 with patch("torch.distributed.get_rank") as rank_method: rank_method.return_value = rank with distributed_checkpoint_dir(step="test") as path: if rank == 0: assert path if rank != 0: assert not os.path.exists(path)