fix: download cifar10 only once per node in cifar10_pytorch example
CNN on Cifar10 from Keras example:
https://github.com/fchollet/keras/blob/master/examples/cifar10_cnn.py
"""
-
+import filelock
+import os
import tempfile
from typing import Any, Dict, Sequence, Tuple, Union, cast
@@ -43,11 +44,11 @@ class CIFARTrial(PyTorchTrial):
def __init__(self, context: PyTorchTrialContext) -> None:
self.context = context
- # Create a unique download directory for each rank so they don't overwrite each
- # other when doing distributed training.
- self.download_directory = tempfile.mkdtemp()
+ self.download_directory = "data"
+ os.makedirs(self.download_directory, exist_ok=True)
- self.model = self.context.wrap_model(nn.Sequential(
+ self.model = self.context.wrap_model(
+ nn.Sequential(
nn.Conv2d(NUM_CHANNELS, IMAGE_SIZE, kernel_size=(3, 3)),
nn.ReLU(),
nn.Conv2d(32, 32, kernel_size=(3, 3)),
@@ -65,14 +66,17 @@ class CIFARTrial(PyTorchTrial):
nn.ReLU(),
nn.Dropout2d(self.context.get_hparam("layer3_dropout")),
nn.Linear(512, NUM_CLASSES),
- ))
+ )
+ )
- self.optimizer = self.context.wrap_optimizer(torch.optim.RMSprop( # type: ignore
+ self.optimizer = self.context.wrap_optimizer(
+ torch.optim.RMSprop( # type: ignore
self.model.parameters(),
lr=self.context.get_hparam("learning_rate"),
weight_decay=self.context.get_hparam("learning_rate_decay"),
alpha=0.9,
- ))
+ )
+ )
def train_batch(
self, batch: TorchData, epoch_idx: int, batch_idx: int
@@ -100,21 +104,22 @@ class CIFARTrial(PyTorchTrial):
accuracy = accuracy_rate(output, labels)
return {"validation_accuracy": accuracy, "validation_error": 1.0 - accuracy}
- def build_training_data_loader(self) -> Any:
+ def _download_dataset(self, train: bool) -> Any:
transform = transforms.Compose(
[transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)
- trainset = torchvision.datasets.CIFAR10(
- root=self.download_directory, train=True, download=True, transform=transform
+ # Use a file lock so that workers on the same node attempt the download one at a time.
+ # The first worker will actually perform the download, while the subsequent workers will
+ # see that the dataset is downloaded and skip.
+ with filelock.FileLock(os.path.join(self.download_directory, "lock")):
+ return torchvision.datasets.CIFAR10(
+ root=self.download_directory, train=train, download=True, transform=transform
)
+
+ def build_training_data_loader(self) -> Any:
+ trainset = self._download_dataset(train=True)
return DataLoader(trainset, batch_size=self.context.get_per_slot_batch_size())
def build_validation_data_loader(self) -> Any:
- transform = transforms.Compose(
- [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
- )
- valset = torchvision.datasets.CIFAR10(
- root=self.download_directory, train=False, download=True, transform=transform
- )
-
+ valset = self._download_dataset(train=False)
return DataLoader(valset, batch_size=self.context.get_per_slot_batch_size())
