add device type in autocast and GradScaler

Signed-off-by: YunLiu <[email protected]>

Author: KumoLiu

automl/DiNTS/search_dints.py

@@ -431,7 +431,7 @@ def main():
     if amp:
         from torch import autocast, GradScaler
 
-        scaler = GradScaler()
+        scaler = GradScaler("cuda")
         if dist.get_rank() == 0:
             print("[info] amp enabled")
 
@@ -487,7 +487,7 @@ def main():
             optimizer.zero_grad()
 
             if amp:
-                with autocast():
+                with autocast("cuda"):
                     outputs = model(inputs)
                     if output_classes == 2:
                         loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
@@ -559,7 +559,7 @@ def main():
             combination_weights = (epoch - num_epochs_warmup) / (num_epochs - num_epochs_warmup)
 
             if amp:
-                with autocast():
+                with autocast("cuda"):
                     outputs_search = model(inputs_search)
                     if output_classes == 2:
                         loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)

automl/DiNTS/train_dints.py

@@ -408,7 +408,7 @@ def main():
     if amp:
         from torch import autocast, GradScaler
 
-        scaler = GradScaler()
+        scaler = GradScaler("cuda")
         if dist.get_rank() == 0:
             print("[info] amp enabled")
 
@@ -450,7 +450,7 @@ def main():
                 param.grad = None
 
             if amp:
-                with autocast():
+                with autocast("cuda"):
                     outputs = model(inputs)
                     if output_classes == 2:
                         loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)

competitions/MICCAI/surgtoolloc/classification_files/train.py

@@ -85,7 +85,7 @@ def main(cfg):
     metric = ConfusionMatrixMetric(metric_name="F1", reduction="mean_batch")
 
     # set other tools
-    scaler = GradScaler()
+    scaler = GradScaler("cuda")
     writer = SummaryWriter(str(cfg.output_dir + f"/fold{cfg.fold}/"))
 
     # train and val loop
@@ -171,11 +171,11 @@ def run_train(
         torch.set_grad_enabled(True)
         if torch.rand(1) > 0.5:
             inputs, labels_a, labels_b, lam = mixup_data(inputs, labels)
-            with autocast():
+            with autocast("cuda"):
                 outputs = model(inputs)
                 loss = lam * loss_function(outputs, labels_a) + (1 - lam) * loss_function(outputs, labels_b)
         else:
-            with autocast():
+            with autocast("cuda"):
                 outputs = model(inputs)
                 loss = loss_function(outputs, labels)
         losses.append(loss.item())

competitions/kaggle/RANZCR/4th_place_solution/train.py

@@ -83,7 +83,7 @@ def main(cfg):
 
     # set other tools
     if cfg.mixed_precision:
-        scaler = GradScaler()
+        scaler = GradScaler("cuda")
     else:
         scaler = None
 
@@ -168,7 +168,7 @@ def run_train(
         torch.set_grad_enabled(True)
 
         if cfg.mixed_precision:
-            with autocast():
+            with autocast("cuda"):
                 output_dict = model(batch)
         else:
             output_dict = model(batch)
@@ -210,7 +210,7 @@ def run_eval(model, val_dataloader, cfg, writer, epoch):
     for batch in val_dataloader:
         batch = cfg.to_device_transform(batch)
         if cfg.mixed_precision:
-            with autocast():
+            with autocast("cuda"):
                 output = model(batch)
         else:
             output = model(batch)
@@ -271,7 +271,7 @@ def run_infer(weights_folder_path, cfg):
             batch = to_device_transform(batch)
             for i, net in enumerate(nets):
                 if cfg.mixed_precision:
-                    with autocast():
+                    with autocast("cuda"):
                         logits = net(batch)["logits"].cpu().numpy()
                 else:
                     logits = net(batch)["logits"].cpu().numpy()

generation/2d_ddpm/2d_ddpm_inpainting.ipynb

@@ -476,7 +476,7 @@
     "    epoch_loss_list = []\n",
     "    val_epoch_loss_list = []\n",
     "\n",
-    "    scaler = GradScaler()\n",
+    "    scaler = GradScaler(\"cuda\")\n",
     "    total_start = time.time()\n",
     "    for epoch in range(max_epochs):\n",
     "        model.train()\n",

generation/2d_ddpm/2d_ddpm_tutorial.ipynb

@@ -494,7 +494,7 @@
     "    epoch_loss_list = []\n",
     "    val_epoch_loss_list = []\n",
     "\n",
-    "    scaler = GradScaler()\n",
+    "    scaler = GradScaler(\"cuda\")\n",
     "    total_start = time.time()\n",
     "    for epoch in range(max_epochs):\n",
     "        model.train()\n",

generation/2d_ddpm/2d_ddpm_tutorial_v_prediction.ipynb

@@ -466,7 +466,7 @@
     "epoch_loss_list = []\n",
     "val_epoch_loss_list = []\n",
     "\n",
-    "scaler = GradScaler()\n",
+    "scaler = GradScaler(\"cuda\")\n",
     "total_start = time.time()\n",
     "for epoch in range(max_epochs):\n",
     "    model.train()\n",

generation/2d_ldm/2d_ldm_tutorial.ipynb

@@ -401,8 +401,8 @@
     "optimizer_d = torch.optim.Adam(discriminator.parameters(), lr=5e-4)\n",
     "\n",
     "# For mixed precision training\n",
-    "scaler_g = GradScaler()\n",
-    "scaler_d = GradScaler()"
+    "scaler_g = GradScaler(\"cuda\")\n",
+    "scaler_d = GradScaler(\"cuda\")"
    ]
   },
   {
@@ -751,7 +751,7 @@
     "val_interval = 40\n",
     "epoch_losses = []\n",
     "val_losses = []\n",
-    "scaler = GradScaler()\n",
+    "scaler = GradScaler(\"cuda\")\n",
     "\n",
     "for epoch in range(max_epochs):\n",
     "    unet.train()\n",

generation/2d_ldm/train_diffusion.py

@@ -182,7 +182,7 @@ def main():
     max_epochs = args.diffusion_train["max_epochs"]
     val_interval = args.diffusion_train["val_interval"]
     autoencoder.eval()
-    scaler = GradScaler()
+    scaler = GradScaler("cuda")
     total_step = 0
     best_val_recon_epoch_loss = 100.0
 

generation/2d_super_resolution/2d_sd_super_resolution.ipynb

@@ -407,8 +407,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "scaler_g = GradScaler()\n",
-    "scaler_d = GradScaler()"
+    "scaler_g = GradScaler(\"cuda\")\n",
+    "scaler_d = GradScaler(\"cuda\")"
    ]
   },
   {
@@ -973,7 +973,7 @@
     "# Optimizers\n",
     "optimizer = torch.optim.Adam(unet.parameters(), lr=5e-5)\n",
     "\n",
-    "scaler_diffusion = GradScaler()\n",
+    "scaler_diffusion = GradScaler(\"cuda\")\n",
     "\n",
     "max_epochs = 200\n",
     "val_interval = 20\n",