accelerate_pytorch_training_loop_with_mixed_precision_and_distributed.py

1import torch
2import torch.nn.functional as F
3from torch.utils.data import DataLoader
4from torchvision import transforms, datasets
5from accelerate import Accelerator
6from torch.optim import Adam
7
8def training_loop():
9    # 1. Initialize the Accelerator
10    accelerator = Accelerator()
11
12    # Define a simple model, optimizer, and data
13    device = accelerator.device
14    model = torch.nn.Linear(784, 10).to(device)
15    optimizer = Adam(model.parameters(), lr=1e-3)
16
17    # Setup data
18    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
19    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
20    train_dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
21
22    # 2. Pass everything to prepare
23    # This handles device placement, distributed sampling, and mixed precision
24    model, optimizer, train_dataloader = accelerator.prepare(
25        model, optimizer, train_dataloader
26    )
27
28    model.train()
29    for epoch in range(1):
30        for batch in train_dataloader:
31            inputs, targets = batch
32            inputs = inputs.view(inputs.size(0), -1)
33
34            optimizer.zero_grad()
35            outputs = model(inputs)
36            loss = F.cross_entropy(outputs, targets)
37
38            # 3. Replace loss.backward() with accelerator.backward(loss)
39            accelerator.backward(loss)
40            
41            optimizer.step()
42
43    print("Training finished!")
44
45if __name__ == "__main__":
46    training_loop()