학습하는데 computational cost가 높아서 어려움이 많은데, amp를 사용하면 더 큰 배치사이즈를 사용할 수 있거나, gpu 스펙에 따라 더 빠른 학습이 가능합니다.

아래와 같이 쉽게 사용할 수 있는데, 간단한 amp 사용 예시 공유합니다.

모두 좋은 결과 있으시길 바래요!

import torch.cuda.amp as amp
  
scaler = amp.GradScaler()
net = YourModel()

# ------------------------
# loss
# ------------------------
loss_fn = nn.L1Loss()

# ------------------------
# Optimizer
# ------------------------
optimizer = optim.Adam(net.parameters(), lr=args.start_lr, weight_decay=args.weight_decay)

for epoch in range(1, args.epochs+1):

    for t, (images, targets) in enumerate(tqdm.tqdm(trainloader)):
      images = images.to(device=device, dtype=torch.float)
      targets = targets.to(device=device, dtype=torch.float)

      net.train()
      optimizer.zero_grad()

      if args.amp:
        with amp.autocast():
          output = net(images)
          # loss
          loss = loss_fn(output, targets)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

      else:
        output = net(images)

        # loss
        loss = loss_fn(output, targets)

        # update
        loss.backward()
        optimizer.step()