Machine Learning & Data Science frameworks - DAWN

We provide a set of pre-populated Conda environments based on the Intel Distribution for Python:

module av intelpython-conda
conda info -e

This module provides environments for PyTorch and Tensorflow.

Please note that Intel code and documentation sometimes refers to ‘XPUs’, a more generic term for accelerators, GPU or otherwise. For Dawn, ‘XPU’ and ‘GPU’ can usually be considered interchangeable.

PyTorch

PyTorch on Intel GPUs is supported by the Intel Extension for PyTorch. On Dawn this version of PyTorch is accessible as a conda environment named pytorch-gpu:

module load intelpython-conda
conda activate pytorch-gpu

Adapting your code to run on the PVCs is straightforward and only takes a few lines of code. For details, see the official documentation - but as a quick example:

import torch
import intel_extension_for_pytorch as ipex

...

# Enable GPU
model = model.to('xpu')
data = data.to('xpu')

model = ipex.optimize(model, dtype=torch.float32)

TensorFlow

Intel supports optimised TensorFlow on both CPU and GPU, using the Intel Extension for TensorFlow. On Dawn this version on TensorFlow is accessible as a conda environment named tensorflow-gpu:

module load intelpython-conda
conda activate tensorflow-gpu

To run on the PVCs, there should be no need to modify your code - the Intel optimised implementation will run automatically on the GPU, assuming it has been installed as intel-extension-for-tensorflow[xpu].

Jax/OpenXLA

Documentation can be found on GitHub: Intel OpenXLA

Julia

This is currently known not to work correctly on PVC GPUs. (Mar 2024)