Skip to content

Add Accelerator API to Imagenet Example#1349

Merged
soumith merged 2 commits into
pytorch:mainfrom
eromomon:eromomon/imagenet-accel
Jun 17, 2025
Merged

Add Accelerator API to Imagenet Example#1349
soumith merged 2 commits into
pytorch:mainfrom
eromomon:eromomon/imagenet-accel

Conversation

@eromomon
Copy link
Copy Markdown
Contributor

Refactor Imagenet example to utilize torch.accelerator API. torch.accelerator API allows to abstract some of the accelerator specifics in the user scripts. By leveraging this API, the code becomes more adaptable to various hardware accelerators.

CC: @msaroufim, @malfet, @dvrogozh

Signed-off-by: eromomon <edgar.romo.montiel@intel.com>
@netlify
Copy link
Copy Markdown

netlify Bot commented May 19, 2025

Deploy Preview for pytorch-examples-preview canceled.

Name Link
🔨 Latest commit cd2b7f2
🔍 Latest deploy log https://app.netlify.com/projects/pytorch-examples-preview/deploys/684cc8192d044d000845933f

Comment thread imagenet/main.py Outdated
if use_accel:
if args.gpu is not None:
torch.accelerator.set_device_index(args.gpu)
print("Use GPU: {} for training".format(args.gpu))
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This print seems off as you don't add it in other cases for example in validate. Is this a leftover debug print? Isn't the print on line 116 not enough?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Solved, print removed

Signed-off-by: Edgar Romo Montiel <edgar.romo.montiel@intel.com>
@eromomon eromomon requested a review from dvrogozh June 14, 2025 00:54
Copy link
Copy Markdown
Contributor

@dvrogozh dvrogozh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@soumith soumith merged commit 37986af into pytorch:main Jun 17, 2025
8 checks passed
Comment thread imagenet/main.py
target = target.cuda(args.gpu, non_blocking=True)
if use_accel:
if args.gpu is not None and device.type=='cuda':
torch.accelerator.set_device_index(argps.gpu)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like a typo here argps instead of args

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wanna send me a patch? I'll merge

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was doing something else on nvidia system and had a spare minute. Here is a PR to fix and typo and add respective test:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants