accelerate_pytorch_training_loop_multi_gpu_tpu_mixed_precision.py

1import torch
2import torch.nn.functional as F
3from torch.utils.data import DataLoader
4from torchvision import transforms, datasets
5from accelerate import Accelerator
6
7def training_loop():
8    # 1. Initialize the Accelerator
9    accelerator = Accelerator()
10
11    # 2. Setup model, optimizer, and data
12    device = accelerator.device
13    model = torch.nn.Sequential(
14        torch.nn.Flatten(),
15        torch.nn.Linear(784, 128),
16        torch.nn.ReLU(),
17        torch.nn.Linear(128, 10)
18    ).to(device)
19    
20    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
21
22    dataset = datasets.MNIST('./data', train=True, download=True, transform=transforms.ToTensor())
23    train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
24
25    # 3. Prepare everything using accelerator.prepare
26    model, optimizer, train_loader = accelerator.prepare(
27        model, optimizer, train_loader
28    )
29
30    model.train()
31    for epoch in range(1):
32        for batch in train_loader:
33            inputs, targets = batch
34            
35            optimizer.zero_grad()
36            outputs = model(inputs)
37            loss = F.cross_entropy(outputs, targets)
38            
39            # 4. Replace loss.backward() with accelerator.backward(loss)
40            accelerator.backward(loss)
41            
42            optimizer.step()
43        
44        print(f"Epoch {epoch} complete.")
45
46if __name__ == "__main__":
47    training_loop()