Add accelerate API support for Super Resolution example (#1358)

Update super_resolution example to support accelerate API

Author: framoncg

run_python_examples.sh

@@ -136,8 +136,8 @@ function fx() {
 }
 
 function super_resolution() {
-  uv run main.py --upscale_factor 3 --batchSize 4 --testBatchSize 100 --nEpochs 1 --lr 0.001 --mps || error "super resolution failed"
-  uv run super_resolve.py --input_image dataset/BSDS300/images/test/16077.jpg --model model_epoch_1.pth --output_filename out.png || error "super resolution upscaling failed"
+  uv run main.py --upscale_factor 3 --batchSize 4 --testBatchSize 100 --nEpochs 1 --lr 0.001 $ACCEL_FLAG || error "super resolution failed"
+  uv run super_resolve.py --input_image dataset/BSDS300/images/test/16077.jpg --model model_epoch_1.pth --output_filename out.png $ACCEL_FLAG || error "super resolution upscaling failed"
 }
 
 function time_sequence_prediction() {

super_resolution/README.md

@@ -5,7 +5,7 @@ This example illustrates how to use the efficient sub-pixel convolution layer de
 ```
 usage: main.py [-h] --upscale_factor UPSCALE_FACTOR [--batchSize BATCHSIZE]
                [--testBatchSize TESTBATCHSIZE] [--nEpochs NEPOCHS] [--lr LR]
-               [--cuda] [--threads THREADS] [--seed SEED]
+               [--accel] [--threads THREADS] [--seed SEED]
 
 PyTorch Super Res Example
 
@@ -16,8 +16,7 @@ optional arguments:
   --testBatchSize       testing batch size
   --nEpochs             number of epochs to train for
   --lr                  Learning Rate. Default=0.01
-  --cuda                use cuda
-  --mps                 enable GPU on macOS
+  --accel                use accelerator
   --threads             number of threads for data loader to use Default=4
   --seed                random seed to use. Default=123
 ```
@@ -29,11 +28,11 @@ This example trains a super-resolution network on the [BSD300 dataset](https://w
 ### Train
 
 ```bash
-python main.py --upscale_factor 3 --batchSize 4 --testBatchSize 100 --nEpochs 30 --lr 0.001
+python main.py --upscale_factor 3 --batchSize 4 --testBatchSize 100 --nEpochs 30 --lr 0.001 --accel
 ```
 
 ### Super Resolve
 
 ```bash
-python super_resolve.py --input_image dataset/BSDS300/images/test/16077.jpg --model model_epoch_30.pth --output_filename out.png
+python super_resolve.py --input_image dataset/BSDS300/images/test/16077.jpg --model model_epoch_30.pth --output_filename out.png --accel
 ```

super_resolution/main.py

@@ -16,26 +16,18 @@
 parser.add_argument('--testBatchSize', type=int, default=10, help='testing batch size')
 parser.add_argument('--nEpochs', type=int, default=2, help='number of epochs to train for')
 parser.add_argument('--lr', type=float, default=0.01, help='Learning Rate. Default=0.01')
-parser.add_argument('--cuda', action='store_true', help='use cuda?')
-parser.add_argument('--mps', action='store_true', default=False, help='enables macOS GPU training')
+parser.add_argument('--accel', action='store_true', help='Enables acceleration for training, if available')
 parser.add_argument('--threads', type=int, default=4, help='number of threads for data loader to use')
 parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')
 opt = parser.parse_args()
 
 print(opt)
 
-if opt.cuda and not torch.cuda.is_available():
-    raise Exception("No GPU found, please run without --cuda")
-if not opt.mps and torch.backends.mps.is_available():
-    raise Exception("Found mps device, please run with --mps to enable macOS GPU")
-
 torch.manual_seed(opt.seed)
-use_mps = opt.mps and torch.backends.mps.is_available()
 
-if opt.cuda:
-    device = torch.device("cuda")
-elif use_mps:
-    device = torch.device("mps")
+
+if opt.accel and torch.accelerator.is_available():
+    device = torch.accelerator.current_accelerator()
 else:
     device = torch.device("cpu")
 

super_resolution/super_resolve.py

@@ -12,7 +12,7 @@
 parser.add_argument('--input_image', type=str, required=True, help='input image to use')
 parser.add_argument('--model', type=str, required=True, help='model file to use')
 parser.add_argument('--output_filename', type=str, help='where to save the output image')
-parser.add_argument('--cuda', action='store_true', help='use cuda')
+parser.add_argument('--accel', action='store_true', help='Enables acceleration device, if available')
 opt = parser.parse_args()
 
 print(opt)
@@ -32,9 +32,10 @@
 img_to_tensor = ToTensor()
 input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])
 
-if opt.cuda:
-    model = model.cuda()
-    input = input.cuda()
+if opt.accel:
+    device = torch.accelerator.current_accelerator()
+    model = model.to(device)
+    input = input.to(device)
 
 out = model(input)
 out = out.cpu()