apply PyTorch AMP

.drone.yml

@@ -20,8 +20,8 @@ steps:
       - pip install -q torchsummary matplotlib pytorch-metric-learning torchinfo torch torchvision tqdm Pillow facenet-pytorch
       - apt update -qq
       - apt install fonts-noto-cjk-extra -y -qq
-      - ls ./
+      # - ls ./
       - mkdir -p data
       - $mount_command
-      - ls data/
+      # - ls data/
       - CI=True python finetune/facenet_transfer_learning.py

finetune/facenet_transfer_learning.py

@@ -158,12 +158,10 @@ for epoch in range(epochs):
                                          label_text=image_folder['train'].classes)
             image_pallets.save(join(save_dir, 'pallets', str(epoch) + '_train.jpg'))
         optimizer.zero_grad()
-        with autocast(dtype=float16):
+        with autocast(dtype=float16,enabled=True):
             images = images.to(device)
             labels = labels.to(device)
-
             outputs = model(images)
-
             loss = criterion(outputs, labels)
         train_loss += loss.item()
 
@@ -187,6 +185,7 @@ for epoch in range(epochs):
                 image_pallets = plot_dataset(dataloader=(images, labels), col_len=6,
                                              label_text=image_folder['train'].classes)
                 image_pallets.save(join(save_dir, 'pallets', str(epoch) + '_val.jpg'))
+            with autocast(dtype=float16, enabled=True):
                 images = images.to(device)
                 labels = labels.to(device)
             outputs = model_gpu(images)