Address review comments: remove cuda parameter from GradScaler, update PyTorch version requirements, revert notebook changes

Co-authored-by: vfdev-5 <[email protected]>

Author: Copilot

examples/cifar10/main.py

@@ -288,7 +288,7 @@ def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, con
     #    - Two progress bars on epochs and optionally on iterations
 
     with_amp = config["with_amp"]
-    scaler = GradScaler('cuda', enabled=with_amp)
+    scaler = GradScaler(enabled=with_amp)
 
     def train_step(engine, batch):
         x, y = batch[0], batch[1]

examples/cifar100_amp_benchmark/benchmark_torch_cuda_amp.py

@@ -25,7 +25,7 @@ def main(dataset_path, batch_size=256, max_epochs=10):
     optimizer = SGD(model.parameters(), lr=0.01)
     criterion = CrossEntropyLoss().to(device)
 
-    scaler = GradScaler('cuda')
+    scaler = GradScaler()
 
     def train_step(engine, batch):
         x = convert_tensor(batch[0], device, non_blocking=True)

examples/cifar10_qat/main.py

@@ -272,7 +272,7 @@ def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, con
     #    - Two progress bars on epochs and optionally on iterations
 
     with_amp = config["with_amp"]
-    scaler = GradScaler('cuda', enabled=with_amp)
+    scaler = GradScaler(enabled=with_amp)
 
     def train_step(engine, batch):
         x, y = batch[0], batch[1]

examples/notebooks/CycleGAN_with_torch_cuda_amp.ipynb

@@ -875,10 +875,10 @@
     "As suggested, we divide the objective by 2 while optimizing D, which slows down the rate at which D learns, relative to the rate of G. \n",
     "\n",
     "According to the paper:\n",
-    "- generator A is trained  minimize $\\text{mean}_{x \\in A}[(D_B(G(x)) \u2212 1)^2]$ and cycle loss $\\text{mean}_{x \\in A}\\left[ |F(G(x)) - x|_1 \\right]$\n",
-    "- generator B is trained  minimize $\\text{mean}_{y \\in B}[(D_A(F(y)) \u2212 1)^2]$ and cycle loss $\\text{mean}_{y \\in B}\\left[ |G(F(y)) - y|_1 \\right]$\n",
-    "- discriminators A is trained to minimize $\\text{mean}_{x \\in A}[(D_A(x) \u2212 1)^2] + \\text{mean}_{y \\in B}[D_A(F(y))^2]$.\n",
-    "- discriminator B is trained to minimize $\\text{mean}_{y \\in B}[(D_B(y) \u2212 1)^2] + \\text{mean}_{x \\in A}[D_B(G(x))^2]$."
+    "- generator A is trained  minimize $\\text{mean}_{x \\in A}[(D_B(G(x)) − 1)^2]$ and cycle loss $\\text{mean}_{x \\in A}\\left[ |F(G(x)) - x|_1 \\right]$\n",
+    "- generator B is trained  minimize $\\text{mean}_{y \\in B}[(D_A(F(y)) − 1)^2]$ and cycle loss $\\text{mean}_{y \\in B}\\left[ |G(F(y)) - y|_1 \\right]$\n",
+    "- discriminators A is trained to minimize $\\text{mean}_{x \\in A}[(D_A(x) − 1)^2] + \\text{mean}_{y \\in B}[D_A(F(y))^2]$.\n",
+    "- discriminator B is trained to minimize $\\text{mean}_{y \\in B}[(D_B(y) − 1)^2] + \\text{mean}_{x \\in A}[D_B(G(x))^2]$."
    ]
   },
   {
@@ -887,7 +887,7 @@
     "id": "JE8dLeEfIl_Z"
    },
    "source": [
-    "We will use [`torch.amp.autocast`](https://pytorch.org/docs/master/amp.html#torch.amp.autocast) and [`torch.amp.GradScaler`](https://pytorch.org/docs/master/amp.html#torch.amp.GradScaler) to perform automatic mixed precision training. Our code follows a [typical mixed precision training example](https://pytorch.org/docs/master/notes/amp_examples.html#typical-mixed-precision-training)."
+    "We will use [`torch.amp.autocast`](https://pytorch.org/docs/master/amp.html#torch.amp.autocast) and [`torch.cuda.amp.GradScaler`](https://pytorch.org/docs/master/amp.html#torch.cuda.amp.GradScaler) to perform automatic mixed precision training. Our code follows a [typical mixed precision training example](https://pytorch.org/docs/master/notes/amp_examples.html#typical-mixed-precision-training)."
    ]
   },
   {
@@ -896,7 +896,7 @@
     "id": "vrJls4p-FRcA"
    },
    "source": [
-    "from torch.amp import GradScaler\n",
+    "from torch.cuda.amp import GradScaler\n",
     "from torch.amp import autocast\n",
     "\n",
     "from ignite.utils import convert_tensor\n",
@@ -924,7 +924,7 @@
     "\n",
     "\n",
     "def compute_loss_discriminator(decision_real, decision_fake):\n",
-    "    # loss = mean (D_b(y) \u2212 1)^2 + mean D_b(G(x))^2    \n",
+    "    # loss = mean (D_b(y) − 1)^2 + mean D_b(G(x))^2    \n",
     "    loss = F.mse_loss(decision_fake, torch.zeros_like(decision_fake))\n",
     "    loss += F.mse_loss(decision_real, torch.ones_like(decision_real))\n",
     "    return loss\n",
@@ -954,10 +954,10 @@
     "        decision_fake_b = discriminator_B(fake_b)\n",
     "\n",
     "        # Compute loss for generators and update generators\n",
-    "        # loss_a2b = GAN loss: mean (D_b(G(x)) \u2212 1)^2 + Forward cycle loss: || F(G(x)) - x ||_1    \n",
+    "        # loss_a2b = GAN loss: mean (D_b(G(x)) − 1)^2 + Forward cycle loss: || F(G(x)) - x ||_1    \n",
     "        loss_a2b = compute_loss_generator(decision_fake_b, real_a, rec_a, lambda_value)    \n",
     "\n",
-    "        # loss_b2a = GAN loss: mean (D_a(F(x)) \u2212 1)^2 + Backward cycle loss: || G(F(y)) - y ||_1\n",
+    "        # loss_b2a = GAN loss: mean (D_a(F(x)) − 1)^2 + Backward cycle loss: || G(F(y)) - y ||_1\n",
     "        loss_b2a = compute_loss_generator(decision_fake_a, real_b, rec_b, lambda_value)\n",
     "\n",
     "        # total generators loss:\n",
@@ -977,10 +977,10 @@
     "        decision_real_a, decision_fake_a = discriminator_forward_pass(discriminator_A, real_a, fake_a.detach(), fake_a_buffer)    \n",
     "        decision_real_b, decision_fake_b = discriminator_forward_pass(discriminator_B, real_b, fake_b.detach(), fake_b_buffer)    \n",
     "        # Compute loss for discriminators and update discriminators\n",
-    "        # loss_a = mean (D_a(y) \u2212 1)^2 + mean D_a(F(x))^2\n",
+    "        # loss_a = mean (D_a(y) − 1)^2 + mean D_a(F(x))^2\n",
     "        loss_a = compute_loss_discriminator(decision_real_a, decision_fake_a)\n",
     "\n",
-    "        # loss_b = mean (D_b(y) \u2212 1)^2 + mean D_b(G(x))^2\n",
+    "        # loss_b = mean (D_b(y) − 1)^2 + mean D_b(G(x))^2\n",
     "        loss_b = compute_loss_discriminator(decision_real_b, decision_fake_b)\n",
     "      \n",
     "        # total discriminators loss:\n",
@@ -1578,4 +1578,4 @@
    "outputs": []
   }
  ]
-}
+}

examples/references/classification/imagenet/main.py

@@ -8,7 +8,7 @@
 try:
     from torch.amp import autocast, GradScaler
 except ImportError:
-    raise RuntimeError("Please, use recent PyTorch version, e.g. >=1.12.0")
+    raise RuntimeError("Please, use recent PyTorch version, e.g. >=2.3.1")
 
 import dataflow as data
 import utils
@@ -139,7 +139,7 @@ def create_trainer(model, optimizer, criterion, train_sampler, config, logger, w
     model_output_transform = config.get("model_output_transform", lambda x: x)
 
     with_amp = config.get("with_amp", True)
-    scaler = GradScaler('cuda', enabled=with_amp)
+    scaler = GradScaler(enabled=with_amp)
 
     def training_step(engine, batch):
         model.train()

examples/references/segmentation/pascal_voc2012/main.py

@@ -8,7 +8,7 @@
 try:
     from torch.amp import autocast, GradScaler
 except ImportError:
-    raise RuntimeError("Please, use recent PyTorch version, e.g. >=1.12.0")
+    raise RuntimeError("Please, use recent PyTorch version, e.g. >=2.3.1")
 
 import dataflow as data
 import utils
@@ -186,7 +186,7 @@ def create_trainer(model, optimizer, criterion, train_sampler, config, logger, w
     model_output_transform = config.get("model_output_transform", lambda x: x)
 
     with_amp = config.get("with_amp", True)
-    scaler = GradScaler('cuda', enabled=with_amp)
+    scaler = GradScaler(enabled=with_amp)
 
     def forward_pass(batch):
         model.train()

examples/transformers/main.py

@@ -297,7 +297,7 @@ def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, con
     #    - Two progress bars on epochs and optionally on iterations
 
     with_amp = config["with_amp"]
-    scaler = GradScaler('cuda', enabled=with_amp)
+    scaler = GradScaler(enabled=with_amp)
 
     def train_step(engine, batch):
         input_batch = batch[0]

ignite/engine/__init__.py

@@ -187,7 +187,7 @@ def supervised_training_step_amp(
     try:
         from torch.amp import autocast, GradScaler
     except ImportError:
-        raise ImportError("Please install torch>=1.12.0 to use amp_mode='amp'.")
+        raise ImportError("Please install torch>=2.3.1 to use amp_mode='amp'.")
 
     if gradient_accumulation_steps <= 0:
         raise ValueError(
@@ -412,8 +412,8 @@ def _check_arg(
             try:
                 from torch.amp import GradScaler
             except ImportError:
-                raise ImportError("Please install torch>=1.6.0 to use scaler argument.")
-            scaler = GradScaler('cuda', enabled=True)
+                raise ImportError("Please install torch>=2.3.1 to use scaler argument.")
+            scaler = GradScaler(enabled=True)
 
     if on_tpu:
         return "tpu", None