search

dusty-nv / jetson-containers

Machine Learning Containers for NVIDIA Jetson and JetPack-L4T
MIT License
1.93k stars 422 forks source link

Improve torchvision version selection logic #416

Open IamShubhamGupto opened 4 months ago

IamShubhamGupto commented 4 months ago

With PyTorch2.0+ being accessible by both Jetpack 5 and Jetpack6 users, the logic to install torchvision must be improved.

The PyTorch - Torchvision compatibility matrix - https://pypi.org/project/torchvision/

The condition for selecting torchvision version should also consider the version of PyTorch already installed -https://github.com/dusty-nv/jetson-containers/blob/0ec423dbed7005de451ac74ed685e635defda05c/packages/pytorch/torchvision/config.py#L6

Im not sure how to go about implementing this as it would require live feedback from the container being built, but id be happy to work on updating this logic

dusty-nv commented 4 months ago

@Shubham what happens is that I set the torchvision version in that config.py to match the default version of pytorch defined for what version of jetpack is being run. If they are out of sync, I need to update them.

I've not thought of an elegant solution yet to runtime versioning of various components like you mentioned - including selecting different CUDA, PyTorch, ect and having the entire stack rebuilt against that. At first I was thinking of using environment variables like CUDA_VERSION and PYTORCH_VERSION that the user could set, but I think that would get cumbersome.

Another idea would be that config.py gets passed the container stack being built so far, which it could then expect - the issue there is that container config is done at initialization, not build time.

Over time I have implemented more versioning of the most important packages, and some now have separate builder containers for wheels that take a long time. Ideally I'd like for these builder containers to push to a self-hosted PyPi index and alleviate the whole thing...

From: Shubham @.> Sent: Monday, March 11, 2024 6:49:20 PM To: dusty-nv/jetson-containers @.> Cc: Subscribed @.***> Subject: [dusty-nv/jetson-containers] Improve torchvision version selection logic (Issue #416)

With PyTorch2.0+ being accessible by both Jetpack 5 and Jetpack6 users, the logic to install torchvision must be improved.

The PyTorch - Torchvision compatibility matrix - https://pypi.org/project/torchvision/

The condition for selecting torchvision version should also consider the version of PyTorch already installed -https://github.com/dusty-nv/jetson-containers/blob/0ec423dbed7005de451ac74ed685e635defda05c/packages/pytorch/torchvision/config.py#L6

Im not sure how to go about implementing this as it would require live feedback from the container being built, but id be happy to work on updating this logic

— Reply to this email directly, view it on GitHubhttps://github.com/dusty-nv/jetson-containers/issues/416, or unsubscribehttps://github.com/notifications/unsubscribe-auth/ADVEGKZCAVEYBOPCRHQKVNTYXY7IBAVCNFSM6AAAAABERFKWH2VHI2DSMVQWIX3LMV43ASLTON2WKOZSGE4DAMZZGQ3TAMA. You are receiving this because you are subscribed to this thread.Message ID: @.***>

IamShubhamGupto commented 4 months ago

Hey @dusty-nv The first idea, while cumbersome, is actually somewhat practical and what I have been doing while building my own customer containers starting from l4t-base as most packages now have assumptions of which CUDA version is being used just by looking at the jetpack version.

im not sure how the config.py would be passed to the container stack but having maybe a build from source flag and version flag could help

Pushing the wheel to PyPi index would be great as I just built the whl for pytorch 2.1 on age Xavier and it ran the whole night