Use torch.accelerator API in mnist examples

mnist/main.py

@@ -82,39 +82,35 @@ def main():
                         help='learning rate (default: 1.0)')
     parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                         help='Learning rate step gamma (default: 0.7)')
-    parser.add_argument('--no-cuda', action='store_true', default=False,
-                        help='disables CUDA training')
-    parser.add_argument('--no-mps', action='store_true', default=False,
-                        help='disables macOS GPU training')
-    parser.add_argument('--dry-run', action='store_true', default=False,
+    parser.add_argument('--no-accel', action='store_true',
+                        help='disables accelerator')
+    parser.add_argument('--dry-run', action='store_true',
                         help='quickly check a single pass')
     parser.add_argument('--seed', type=int, default=1, metavar='S',
                         help='random seed (default: 1)')
     parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                         help='how many batches to wait before logging training status')
-    parser.add_argument('--save-model', action='store_true', default=False,
+    parser.add_argument('--save-model', action='store_true', 
                         help='For Saving the current Model')
     args = parser.parse_args()
-    use_cuda = not args.no_cuda and torch.cuda.is_available()
-    use_mps = not args.no_mps and torch.backends.mps.is_available()
+
+    use_accel = not args.no_accel and torch.accelerator.is_available()
 
     torch.manual_seed(args.seed)
 
-    if use_cuda:
-        device = torch.device("cuda")
-    elif use_mps:
-        device = torch.device("mps")
+    if use_accel:
+        device = torch.accelerator.current_accelerator()
     else:
         device = torch.device("cpu")
 
     train_kwargs = {'batch_size': args.batch_size}
     test_kwargs = {'batch_size': args.test_batch_size}
-    if use_cuda:
-        cuda_kwargs = {'num_workers': 1,
+    if use_accel:
+        accel_kwargs = {'num_workers': 1,
                        'pin_memory': True,
                        'shuffle': True}
-        train_kwargs.update(cuda_kwargs)
-        test_kwargs.update(cuda_kwargs)
+        train_kwargs.update(accel_kwargs)
+        test_kwargs.update(accel_kwargs)
 
     transform=transforms.Compose([
         transforms.ToTensor(),

mnist/requirements.txt

@@ -1,2 +1,2 @@
 torch
-torchvision==0.20.0
+torchvision

mnist_forward_forward/README.md

@@ -16,8 +16,7 @@ optional arguments:
   -h, --help            show this help message and exit
   --epochs EPOCHS       number of epochs to train (default: 1000)
   --lr LR               learning rate (default: 0.03)
-  --no_cuda             disables CUDA training
-  --no_mps              disables MPS training
+  --no_accel            disables accelerator
   --seed SEED           random seed (default: 1)
   --save_model          For saving the current Model
   --train_size TRAIN_SIZE

mnist_forward_forward/main.py

@@ -102,18 +102,14 @@ def train(self, x_pos, x_neg):
         help="learning rate (default: 0.03)",
     )
     parser.add_argument(
-        "--no_cuda", action="store_true", default=False, help="disables CUDA training"
-    )
-    parser.add_argument(
-        "--no_mps", action="store_true", default=False, help="disables MPS training"
+        "--no_accel", action="store_true", help="disables accelerator"
     )
     parser.add_argument(
         "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)"
     )
     parser.add_argument(
         "--save_model",
         action="store_true",
-        default=False,
         help="For saving the current Model",
     )
     parser.add_argument(
@@ -126,7 +122,6 @@ def train(self, x_pos, x_neg):
     parser.add_argument(
         "--save-model",
         action="store_true",
-        default=False,
         help="For Saving the current Model",
     )
     parser.add_argument(
@@ -137,22 +132,19 @@ def train(self, x_pos, x_neg):
         help="how many batches to wait before logging training status",
     )
     args = parser.parse_args()
-    use_cuda = not args.no_cuda and torch.cuda.is_available()
-    use_mps = not args.no_mps and torch.backends.mps.is_available()
-    if use_cuda:
-        device = torch.device("cuda")
-    elif use_mps:
-        device = torch.device("mps")
+    use_accel = not args.no_accel and torch.accelerator.is_available()
+    if use_accel:
+        device = torch.accelerator.current_accelerator()
     else:
         device = torch.device("cpu")
 
     train_kwargs = {"batch_size": args.train_size}
     test_kwargs = {"batch_size": args.test_size}
 
-    if use_cuda:
-        cuda_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True}
-        train_kwargs.update(cuda_kwargs)
-        test_kwargs.update(cuda_kwargs)
+    if use_accel:
+        accel_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True}
+        train_kwargs.update(accel_kwargs)
+        test_kwargs.update(accel_kwargs)
 
     transform = Compose(
         [

mnist_forward_forward/requirements.txt

@@ -1,2 +1,2 @@
 torch
-torchvision==0.20.0
+torchvision

mnist_rnn/README.md

@@ -8,3 +8,18 @@ pip install -r requirements.txt
 python main.py
 # CUDA_VISIBLE_DEVICES=2 python main.py  # to specify GPU id to ex. 2
 ```
+
+```bash
+optional arguments:
+  -h, --help            show this help message and exit
+  --batch_size          input batch_size for training (default:64)
+  --testing_batch_size  input batch size for testing (default: 1000)
+  --epochs EPOCHS       number of epochs to train (default: 14)
+  --lr LR               learning rate (default: 0.1)
+  --gamma               learning rate step gamma (default: 0.7)
+  --accel               enables accelerator
+  --seed SEED           random seed (default: 1)
+  --save_model          For saving the current Model
+  --log_interval        how many batches to wait before logging training status
+  --dry-run             quickly check a single pass
+```

mnist_rnn/main.py

@@ -91,32 +91,26 @@ def main():
                         help='learning rate (default: 0.1)')
     parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                         help='learning rate step gamma (default: 0.7)')
-    parser.add_argument('--cuda', action='store_true', default=False,
-                        help='enables CUDA training')
-    parser.add_argument('--mps', action="store_true", default=False,
-                        help="enables MPS training")
-    parser.add_argument('--dry-run', action='store_true', default=False,
+    parser.add_argument('--accel', action='store_true',
+                        help='enables accelerator')
+    parser.add_argument('--dry-run', action='store_true',
                         help='quickly check a single pass')
     parser.add_argument('--seed', type=int, default=1, metavar='S',
                         help='random seed (default: 1)')
     parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                         help='how many batches to wait before logging training status')
-    parser.add_argument('--save-model', action='store_true', default=False,
+    parser.add_argument('--save-model', action='store_true',
                         help='for Saving the current Model')
     args = parser.parse_args()
 
-    if args.cuda and not args.mps:
-        device = "cuda"
-    elif args.mps and not args.cuda:
-        device = "mps"
+    if args.accel:
+        device = torch.accelerator.current_accelerator()
     else:
-        device = "cpu"
-
-    device = torch.device(device)
+        device = torch.device("cpu")
 
     torch.manual_seed(args.seed)
 
-    kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
+    kwargs = {'num_workers': 1, 'pin_memory': True} if args.accel else {}
     train_loader = torch.utils.data.DataLoader(
         datasets.MNIST('../data', train=True, download=True,
                        transform=transforms.Compose([

mnist_rnn/requirements.txt

@@ -1,2 +1,2 @@
 torch
-torchvision==0.20.0
+torchvision

run_python_examples.sh

@@ -93,7 +93,7 @@ function mnist() {
   uv run main.py --epochs 1 --dry-run || error "mnist example failed"
 }
 function mnist_forward_forward() {
-  uv run main.py --epochs 1 --no_mps --no_cuda || error "mnist forward forward failed"
+  uv run main.py --epochs 1 --no_accel || error "mnist forward forward failed"
 
 }
 function mnist_hogwild() {