Without AMP
import torch
from torch.cuda.amp import autocast, GradScaler
from torchvision import models
model = models.mobilenet_v2(pretrained=True).cuda()
loss_fnc = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
X = torch.randn((32,3,300,300), dtype=torch.float32).cuda()
y = torch.randint(0, 1000, (32,), dtype=torch.long).cuda()
model.train()
for j in range(30):
optimizer.zero_grad()
y_hat = model(X)
loss = loss_fnc(y_hat, y)
loss.backward()
optimizer.step()
print (loss.item())
Output:
8.039933204650879
5.690041542053223
3.4787116050720215
1.607206106185913
0.6231755614280701
0.23825135827064514
0.08544095605611801
0.04335329309105873
0.016259444877505302
0.01174827478826046
0.0069425650872290134
0.004459714516997337
0.003734807949513197
0.0024659112095832825
0.0027059323620051146
........................
But with AMP it gives nan’s
model = models.mobilenet_v2(pretrained=True).cuda()
loss_fnc = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
X = torch.randn((32,3,300,300), dtype=torch.float32).cuda()
y = torch.randint(0, 1000, (32,), dtype=torch.long).cuda()
scaler = GradScaler()
model.train()
for j in range(30):
optimizer.zero_grad()
with autocast():
y_hat = model(X)
loss = loss_fnc(y_hat, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print (loss.item())
Output:
8.393239974975586
nan
nan
nan
nan
nan
nan
nan
nan
nan
nan
nan
nan
nan
nan
.......................
