diff --git a/run_python_examples.sh b/run_python_examples.sh index cebd2aec10..c9c90f9552 100755 --- a/run_python_examples.sh +++ b/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() { diff --git a/super_resolution/README.md b/super_resolution/README.md index b21a8c4af4..68bbd473af 100644 --- a/super_resolution/README.md +++ b/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 ``` diff --git a/super_resolution/main.py b/super_resolution/main.py index 8c5519582a..1112f52d17 100644 --- a/super_resolution/main.py +++ b/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") diff --git a/super_resolution/super_resolve.py b/super_resolution/super_resolve.py index 72ca24f022..d21ae9d0fa 100644 --- a/super_resolution/super_resolve.py +++ b/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()