basic_usage.py

"""
Minimal example demonstrating hyrax usage.
"""

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset
from hyrax import DistributedTrainer


class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 1)
    
    def forward(self, x):
        return self.layer(x)


# create dummy datasets
def create_dummy_dataset(seed):
    torch.manual_seed(seed)
    x = torch.randn(1000, 10)
    y = torch.randn(1000, 1)
    return TensorDataset(x, y)


# create three datasets
datasets = [
    create_dummy_dataset(i) 
    for i in range(3)
]

# train
trainer = DistributedTrainer(
    model=TinyModel,
    datasets=datasets
)

results = trainer.train(epochs=20)

print("\nresults:")
for r in results:
    print(f"worker {r['worker_id']}: loss = {r['final_loss']:.4f}")