File size: 991 Bytes
309b968 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """Template: copy this, fill in your model + data, then set
DAISY_TASK=my_task_template:MyTask
(make sure the file is importable, e.g. run daisychain-train from this folder or
pip-install your package). Keep build_model deterministic so every node starts
from identical weights.
"""
import torch
import torch.nn as nn
class MyTask:
def build_model(self) -> nn.Module:
torch.manual_seed(0) # identical init on every node
# TODO: return YOUR model
return nn.Sequential(nn.Linear(16, 64), nn.ReLU(), nn.Linear(64, 10))
def sample(self, n: int):
# TODO: return n training samples from THIS node's data shard as (X, y).
# For real data, shard by rank (e.g. different files/rows per RANK).
X = torch.randn(n, 16)
y = torch.randint(0, 10, (n,))
return X, y
def loss(self, model, X, y):
# TODO: your loss (mean over the batch)
return nn.functional.cross_entropy(model(X), y)
|