55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
|
"""MNIST job module.
|
||
|
|
|
||
|
|
Exports make_model() and make_dataloader() for the scheduler + worker.
|
||
|
|
Falls back to synthetic data if torchvision is not installed.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
from torch.utils.data import DataLoader, Subset, TensorDataset
|
||
|
|
|
||
|
|
|
||
|
|
class MNISTNet(nn.Module):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.fc1 = nn.Linear(784, 128)
|
||
|
|
self.fc2 = nn.Linear(128, 10)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x = x.view(x.size(0), -1)
|
||
|
|
x = F.relu(self.fc1(x))
|
||
|
|
return self.fc2(x)
|
||
|
|
|
||
|
|
|
||
|
|
def make_model() -> nn.Module:
|
||
|
|
return MNISTNet()
|
||
|
|
|
||
|
|
|
||
|
|
def make_dataloader(
|
||
|
|
rank: int, world_size: int, batch_size: int = 64
|
||
|
|
) -> DataLoader:
|
||
|
|
try:
|
||
|
|
from torchvision import datasets, transforms
|
||
|
|
|
||
|
|
dataset = datasets.MNIST(
|
||
|
|
"/tmp/mnist_data",
|
||
|
|
train=True,
|
||
|
|
download=True,
|
||
|
|
transform=transforms.ToTensor(),
|
||
|
|
)
|
||
|
|
except (ImportError, Exception):
|
||
|
|
# synthetic fallback — same shape as MNIST
|
||
|
|
n = 4096
|
||
|
|
x = torch.randn(n, 1, 28, 28)
|
||
|
|
y = torch.randint(0, 10, (n,))
|
||
|
|
dataset = TensorDataset(x, y)
|
||
|
|
|
||
|
|
# shard by rank: interleaved assignment
|
||
|
|
indices = list(range(rank, len(dataset), world_size))
|
||
|
|
subset = Subset(dataset, indices)
|
||
|
|
|
||
|
|
return DataLoader(subset, batch_size=batch_size, shuffle=True, drop_last=True)
|