Use torch.accelerator API in Imagenet example

Author: eromomon

imagenet/README.md

@@ -33,9 +33,7 @@ python main.py -a resnet18 --dummy
 
 ## Multi-processing Distributed Data Parallel Training
 
-If running on CUDA, you should always use the NCCL backend for multi-processing distributed training since it currently provides the best distributed training performance.
-
-For XPU multiprocessing is not supported as of PyTorch 2.6.
+You should always use the NCCL backend for multi-processing distributed training since it currently provides the best distributed training performance.
 
 ### Single node, multiple GPUs:
 
@@ -61,7 +59,7 @@ python main.py -a resnet50 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' --dist-backen
 
 ```bash
 usage: main.py [-h] [-a ARCH] [-j N] [--epochs N] [--start-epoch N] [-b N] [--lr LR] [--momentum M] [--wd W] [-p N] [--resume PATH] [-e] [--pretrained] [--world-size WORLD_SIZE] [--rank RANK]
-               [--dist-url DIST_URL] [--dist-backend DIST_BACKEND] [--seed SEED] [--gpu GPU] [--multiprocessing-distributed] [--dummy]
+               [--dist-url DIST_URL] [--dist-backend DIST_BACKEND] [--seed SEED] [--gpu GPU] [--no-accel][--multiprocessing-distributed] [--dummy]
                [DIR]
 
 PyTorch ImageNet Training
@@ -98,6 +96,7 @@ optional arguments:
                         distributed backend
   --seed SEED           seed for initializing training.
   --gpu GPU             GPU id to use.
+  --no-accel            disables accelerator
   --multiprocessing-distributed
                         Use multi-processing distributed training to launch N processes per node, which has N GPUs. This is the fastest way to use PyTorch for either single node or multi node data parallel
                         training

imagenet/main.py

@@ -71,6 +71,8 @@
                     help='seed for initializing training. ')
 parser.add_argument('--gpu', default=None, type=int,
                     help='GPU id to use.')
+parser.add_argument('--no-accel', action='store_true',
+                    help='disables accelerator')
 parser.add_argument('--multiprocessing-distributed', action='store_true',
                     help='Use multi-processing distributed training to launch '
                          'N processes per node, which has N GPUs. This is the '
@@ -104,8 +106,17 @@ def main():
 
     args.distributed = args.world_size > 1 or args.multiprocessing_distributed
 
-    if torch.cuda.is_available():
-        ngpus_per_node = torch.cuda.device_count()
+    use_accel = not args.no_accel and torch.accelerator.is_available()
+
+    if use_accel:
+        device = torch.accelerator.current_accelerator()
+    else:
+        device = torch.device("cpu")
+
+    print(f"Using device: {device}")
+
+    if device.type =='cuda':
+        ngpus_per_node = torch.accelerator.device_count()
         if ngpus_per_node == 1 and args.dist_backend == "nccl":
             warnings.warn("nccl backend >=2.5 requires GPU count>1, see https://github.com/NVIDIA/nccl/issues/103 perhaps use 'gloo'")
     else:
@@ -127,6 +138,13 @@ def main_worker(gpu, ngpus_per_node, args):
     global best_acc1
     args.gpu = gpu
 
+    use_accel = not args.no_accel and torch.accelerator.is_available()
+
+    if use_accel:
+        device = torch.accelerator.current_accelerator()
+    else:
+        device = torch.device("cpu")
+
     if args.gpu is not None:
         print("Use GPU: {} for training".format(args.gpu))
 
@@ -147,16 +165,16 @@ def main_worker(gpu, ngpus_per_node, args):
         print("=> creating model '{}'".format(args.arch))
         model = models.__dict__[args.arch]()
 
-    if not torch.cuda.is_available() and not torch.backends.mps.is_available() and not torch.xpu.is_available():
+    if not use_accel:
         print('using CPU, this will be slow')
     elif args.distributed:
         # For multiprocessing distributed, DistributedDataParallel constructor
         # should always set the single device scope, otherwise,
         # DistributedDataParallel will use all available devices.
-        if torch.cuda.is_available():
+        if device.type == 'cuda':
             if args.gpu is not None:
-                torch.cuda.set_device(args.gpu)
-                model.cuda(args.gpu)
+                torch.accelerator.set_device_index(args.gpu)
+                model.to(device)
                 # When using a single GPU per process and per
                 # DistributedDataParallel, we need to divide the batch size
                 # ourselves based on the total number of GPUs of the current node.
@@ -168,15 +186,12 @@ def main_worker(gpu, ngpus_per_node, args):
                 # DistributedDataParallel will divide and allocate batch_size to all
                 # available GPUs if device_ids are not set
                 model = torch.nn.parallel.DistributedDataParallel(model)
-    elif args.gpu is not None and torch.cuda.is_available():
-        torch.cuda.set_device(args.gpu)
-        model = model.cuda(args.gpu)
-    elif torch.xpu.is_available():
-        device = torch.device("xpu")
-        model = model.to(device)
-    elif torch.backends.mps.is_available():
-        device = torch.device("mps")
-        model = model.to(device)
+
+    elif args.gpu is not None and device.type=='cuda':
+        torch.accelerator.set_device_index(args.gpu)
+        model.to(device)
+    elif device.type != 'cuda':
+        model.to(device)
     else:
         # DataParallel will divide and allocate batch_size to all available GPUs
         if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
@@ -185,20 +200,6 @@ def main_worker(gpu, ngpus_per_node, args):
         else:
             model = torch.nn.DataParallel(model).cuda()
 
-    if torch.cuda.is_available():
-        if args.gpu:
-            device = torch.device('cuda:{}'.format(args.gpu))
-        else:
-            device = torch.device("cuda")
-    elif torch.xpu.is_available():
-        device = torch.device("xpu")
-    elif torch.backends.mps.is_available():
-        device = torch.device("mps")
-    else:
-        device = torch.device("cpu")
-
-    print (f"Device to use: ", {device.type})
-
     # define loss function (criterion), optimizer, and learning rate scheduler
     criterion = nn.CrossEntropyLoss().to(device)
 
@@ -215,7 +216,7 @@ def main_worker(gpu, ngpus_per_node, args):
             print("=> loading checkpoint '{}'".format(args.resume))
             if args.gpu is None:
                 checkpoint = torch.load(args.resume)
-            elif torch.cuda.is_available():
+            elif device.type=='cuda':
                 # Map model to be loaded to specified single gpu.
                 loc = 'cuda:{}'.format(args.gpu)
                 checkpoint = torch.load(args.resume, map_location=loc)
@@ -310,11 +311,14 @@ def main_worker(gpu, ngpus_per_node, args):
 
 
 def train(train_loader, model, criterion, optimizer, epoch, device, args):
-    batch_time = AverageMeter('Time', ':6.3f')
-    data_time = AverageMeter('Data', ':6.3f')
-    losses = AverageMeter('Loss', ':.4e')
-    top1 = AverageMeter('Acc@1', ':6.2f')
-    top5 = AverageMeter('Acc@5', ':6.2f')
+    
+    use_accel = not args.no_accel and torch.accelerator.is_available()
+
+    batch_time = AverageMeter('Time', use_accel, ':6.3f', Summary.NONE)
+    data_time = AverageMeter('Data', use_accel, ':6.3f', Summary.NONE)
+    losses = AverageMeter('Loss', use_accel, ':.4e', Summary.NONE)
+    top1 = AverageMeter('Acc@1', use_accel, ':6.2f', Summary.NONE)
+    top5 = AverageMeter('Acc@5', use_accel, ':6.2f', Summary.NONE)
     progress = ProgressMeter(
         len(train_loader),
         [batch_time, data_time, losses, top1, top5],
@@ -357,23 +361,27 @@ def train(train_loader, model, criterion, optimizer, epoch, device, args):
 
 def validate(val_loader, model, criterion, args):
 
+    use_accel = not args.no_accel and torch.accelerator.is_available()
+
     def run_validate(loader, base_progress=0):
+
+        if use_accel:
+            device = torch.accelerator.current_accelerator()
+        else:
+            device = torch.device("cpu")
+
         with torch.no_grad():
             end = time.time()
             for i, (images, target) in enumerate(loader):
                 i = base_progress + i
-
-                if torch.cuda.is_available():
-                    if args.gpu is not None:
+                if use_accel:
+                    if args.gpu is not None and device.type=='cuda':
+                        torch.accelerator.set_device_index(argps.gpu)
                         images = images.cuda(args.gpu, non_blocking=True)
-                    target = target.cuda(args.gpu, non_blocking=True)
-
-                elif torch.xpu.is_available():
-                    images = images.to("xpu")
-                    target = target.to("xpu")
-                elif torch.backends.mps.is_available():
-                    images = images.to('mps')
-                    target = target.to('mps')
+                        target = target.cuda(args.gpu, non_blocking=True)
+                    else:
+                        images = images.to(device)
+                        target = target.to(device)
 
                 # compute output
                 output = model(images)
@@ -392,10 +400,10 @@ def run_validate(loader, base_progress=0):
                 if i % args.print_freq == 0:
                     progress.display(i + 1)
 
-    batch_time = AverageMeter('Time', ':6.3f', Summary.NONE)
-    losses = AverageMeter('Loss', ':.4e', Summary.NONE)
-    top1 = AverageMeter('Acc@1', ':6.2f', Summary.AVERAGE)
-    top5 = AverageMeter('Acc@5', ':6.2f', Summary.AVERAGE)
+    batch_time = AverageMeter('Time', use_accel, ':6.3f', Summary.NONE)
+    losses = AverageMeter('Loss', use_accel, ':.4e', Summary.NONE)
+    top1 = AverageMeter('Acc@1', use_accel, ':6.2f', Summary.AVERAGE)
+    top5 = AverageMeter('Acc@5', use_accel, ':6.2f', Summary.AVERAGE)
     progress = ProgressMeter(
         len(val_loader) + (args.distributed and (len(val_loader.sampler) * args.world_size < len(val_loader.dataset))),
         [batch_time, losses, top1, top5],
@@ -435,8 +443,9 @@ class Summary(Enum):
 
 class AverageMeter(object):
     """Computes and stores the average and current value"""
-    def __init__(self, name, fmt=':f', summary_type=Summary.AVERAGE):
+    def __init__(self, name, use_accel, fmt=':f', summary_type=Summary.AVERAGE):
         self.name = name
+        self.use_accel = use_accel
         self.fmt = fmt
         self.summary_type = summary_type
         self.reset()
@@ -453,13 +462,9 @@ def update(self, val, n=1):
         self.count += n
         self.avg = self.sum / self.count
 
-    def all_reduce(self):
-        if torch.cuda.is_available():
-            device = torch.device("cuda")
-        elif torch.xpu.is_available():
-            device = torch.device("xpu")        
-        elif torch.backends.mps.is_available():
-            device = torch.device("mps")
+    def all_reduce(self):    
+        if use_accel:
+            device = torch.accelerator.current_accelerator()
         else:
             device = torch.device("cpu")
         total = torch.tensor([self.sum, self.count], dtype=torch.float32, device=device)

imagenet/requirements.txt

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