Set up conda environment for ROCm and PyTorch
Sets up a conda environment with PyTorch optimized for ROCm GPUs. Use this when configuring AMD GPU workstations for deep learning to ensure proper driver and library compatibility.
/plugin marketplace add danielrosehill/linux-desktop-plugin/plugin install lan-manager@danielrosehillYou are helping the user set up a conda environment optimized for ROCm and PyTorch.
Check if conda is installed
conda --versionwget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.shVerify ROCm is available on system
rocminforocminfo | grep "Name:" | head -1Create conda environment
conda create -n rocm-pytorch python=3.11 -y
conda activate rocm-pytorch
Install PyTorch with ROCm support
# For ROCm 6.0
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.0
# For ROCm 5.7
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.7
Install essential ML libraries
conda install -c conda-forge numpy scipy matplotlib jupyter ipython -y
pip install pandas scikit-learn
Install deep learning tools
pip install transformers accelerate datasets
pip install tensorboard
pip install onnx onnxruntime
Test PyTorch ROCm integration
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}") # ROCm uses CUDA API
if torch.cuda.is_available():
print(f"Device name: {torch.cuda.get_device_name(0)}")
print(f"Device count: {torch.cuda.device_count()}")
Create activation script
~/scripts/activate-rocm-pytorch.sh:
#!/bin/bash
eval "$(conda shell.bash hook)"
conda activate rocm-pytorch
echo "ROCm PyTorch environment activated"
python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA available: {torch.cuda.is_available()}')"
Optional: Install additional tools
timm - PyTorch image modelstorchmetrics - Metricslightning - PyTorch Lightningeinops - Tensor operationsProvide a summary showing: