wehub-resource-sync
79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
import os
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn
|
|
from torch.utils.data import DataLoader
|
|
from torchvision import datasets
|
|
from torchvision.transforms import transforms
|
|
|
|
import ray
|
|
from ray.train import ScalingConfig
|
|
from ray.train.examples.horovod.horovod_pytorch_example import (
|
|
Net,
|
|
train_func as hvd_train_func,
|
|
)
|
|
from ray.train.horovod import HorovodTrainer
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
def run_image_prediction(model: torch.nn.Module, images: torch.Tensor) -> torch.Tensor:
|
|
model.eval()
|
|
with torch.no_grad():
|
|
return torch.exp(model(images)).argmax(dim=1)
|
|
|
|
|
|
def test_horovod(ray_start_4_cpus):
|
|
def train_func(config):
|
|
result = hvd_train_func(config)
|
|
assert len(result) == epochs
|
|
assert result[-1] < result[0]
|
|
|
|
num_workers = 1
|
|
epochs = 10
|
|
scaling_config = ScalingConfig(num_workers=num_workers)
|
|
config = {"num_epochs": epochs, "save_model_as_dict": False}
|
|
trainer = HorovodTrainer(
|
|
train_loop_per_worker=train_func,
|
|
train_loop_config=config,
|
|
scaling_config=scaling_config,
|
|
)
|
|
result = trainer.fit()
|
|
|
|
model = Net()
|
|
with result.checkpoint.as_directory() as checkpoint_dir:
|
|
model.load_state_dict(torch.load(os.path.join(checkpoint_dir, "model.pt")))
|
|
|
|
# Find some test data to run on.
|
|
test_set = datasets.MNIST(
|
|
"./data",
|
|
train=False,
|
|
download=True,
|
|
transform=transforms.Compose(
|
|
[transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
|
|
),
|
|
)
|
|
|
|
test_dataloader = DataLoader(test_set, batch_size=10)
|
|
test_dataloader_iter = iter(test_dataloader)
|
|
images, labels = next(
|
|
test_dataloader_iter
|
|
) # only running a batch inference of 10 images
|
|
predicted_labels = run_image_prediction(model, images)
|
|
assert torch.equal(predicted_labels, labels)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|