pytorch-extra

Version:: 1.9.1
Category:: ai
Cluster:: Loki

Author / Distributor

Description

pytorch-extra provides PyTorch 1.9.1 pre-configured with additional deep learning utilities and dependencies to simplify model development and training workflows.

Documentation

import torch
import torchvision
import torchaudio

print(torch.__version__)           # Should output 1.9.1
print(torchvision.__version__)     # torchvision version
print(torch.cuda.is_available())   # Check for CUDA support

# Example - Load a pretrained ResNet18
model = torchvision.models.resnet18(pretrained=True)
model.eval()

# Run dummy inference
x = torch.randn(1, 3, 224, 224)
y = model(x)

Examples/Usage

Load a module (Python 3.9 + CUDA 11.2):

$ module load pytorch-extra-py39-cuda11.2-gcc9/1.9.1

Launch Jupyter Notebook:

$ jupyter-notebook

Quick GPU test in Python:

import torch
print("CUDA device:", torch.cuda.get_device_name(0))

Visualize image data:

import torchvision.transforms as T
from PIL import Image
import matplotlib.pyplot as plt

img = Image.open("cat.jpg")
transform = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])
plt.imshow(transform(img).permute(1, 2, 0))
plt.show()

Unload the module:

$ module unload pytorch-extra-py39-cuda11.2-gcc9/1.9.1

Installation

Source: https://github.com/pytorch/pytorch/releases/tag/v1.9.1