From f677078c04c386f441ce3b1dcb0b4b4cbbdcfb48 Mon Sep 17 00:00:00 2001 From: eromomon Date: Wed, 30 Apr 2025 15:33:36 -0700 Subject: [PATCH 1/2] Use torch.acceleratort API in VAE example --- vae/README.md | 3 +-- vae/main.py | 25 ++++++++++++++----------- vae/requirements.txt | 2 +- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/vae/README.md b/vae/README.md index cda6a33672..9a8d76f2f9 100644 --- a/vae/README.md +++ b/vae/README.md @@ -14,8 +14,7 @@ The main.py script accepts the following arguments: optional arguments: --batch-size input batch size for training (default: 128) --epochs number of epochs to train (default: 10) - --no-cuda enables CUDA training - --mps enables GPU on macOS + --accel use accelerator --seed random seed (default: 1) --log-interval how many batches to wait before logging training status ``` \ No newline at end of file diff --git a/vae/main.py b/vae/main.py index d69833fbe0..9a6850ccd1 100644 --- a/vae/main.py +++ b/vae/main.py @@ -13,28 +13,31 @@ help='input batch size for training (default: 128)') parser.add_argument('--epochs', type=int, default=10, metavar='N', help='number of epochs to train (default: 10)') -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('--accel', action='store_true', + help='use accelerator') 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') args = parser.parse_args() -args.cuda = not args.no_cuda and torch.cuda.is_available() -use_mps = not args.no_mps and torch.backends.mps.is_available() + torch.manual_seed(args.seed) -if args.cuda: - device = torch.device("cuda") -elif use_mps: - device = torch.device("mps") +if args.accel and not torch.accelerator.is_available(): + print("ERROR: accelerator is not available, try running on CPU") + sys.exit(1) +if not args.accel and torch.accelerator.is_available(): + print("WARNING: accelerator is available, run with --accel to enable it") + +if args.accel: + device = torch.accelerator.current_accelerator() else: device = torch.device("cpu") -kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} +print(f"Using device: {device}") + +kwargs = {'num_workers': 1, 'pin_memory': True} if device=="cuda" else {} train_loader = torch.utils.data.DataLoader( datasets.MNIST('../data', train=True, download=True, transform=transforms.ToTensor()), diff --git a/vae/requirements.txt b/vae/requirements.txt index 9a7236577d..73348074bf 100644 --- a/vae/requirements.txt +++ b/vae/requirements.txt @@ -1,4 +1,4 @@ torch -torchvision==0.20.0 +torchvision tqdm six From c0bb4fdba48c3a79ecbc18b94546ece54a1d0907 Mon Sep 17 00:00:00 2001 From: eromomon Date: Thu, 1 May 2025 15:40:26 -0700 Subject: [PATCH 2/2] Use torch.accelerator API in VAE examples, fix README --- vae/README.md | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/vae/README.md b/vae/README.md index 9a8d76f2f9..81e6458d7a 100644 --- a/vae/README.md +++ b/vae/README.md @@ -8,13 +8,12 @@ pip install -r requirements.txt python main.py ``` -The main.py script accepts the following arguments: +The main.py script accepts the following optional arguments: ```bash -optional arguments: - --batch-size input batch size for training (default: 128) - --epochs number of epochs to train (default: 10) - --accel use accelerator - --seed random seed (default: 1) - --log-interval how many batches to wait before logging training status +--batch-size input batch size for training (default: 128) +--epochs number of epochs to train (default: 10) +--accel use accelerator +--seed random seed (default: 1) +--log-interval how many batches to wait before logging training status ``` \ No newline at end of file