Set Up a GPU-Accelerated Python Environment on Ubuntu
This guide sets up a Conda-based Python environment for GPU-accelerated deep learning on Ubuntu, installing TensorFlow, PyTorch, and common data science libraries.
Prerequisites
- Ubuntu server with an NVIDIA GPU
- NVIDIA drivers installed (
nvidia-smishould show the GPU) - Miniconda installed (see below if not yet installed)
Install Miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
bash miniconda.sh
conda config --set auto_activate_base false
Install CUDA Toolkit
apt-get install nvidia-cuda-toolkit
Create the Environment and Install Packages
# Create and activate environment
conda create -n aiml_env python=3.10 -y
conda activate aiml_env
# Install TensorFlow with CUDA support
conda install tensorflow[and-cuda]
# Install data science packages
conda install numpy pandas matplotlib scikit-learn jupyter
# Install PyTorch with CUDA 11.8
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 \
-c pytorch -c nvidia
Verify GPU Access
# Check GPU is visible
lspci | grep -i nvidia
# Verify PyTorch sees the GPU
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"
Launch Jupyter Notebook
conda activate aiml_env
jupyter notebook --ip=0.0.0.0 --no-browser
Notes
- Accept the Anaconda terms of service if prompted:
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main - Match the
pytorch-cudaversion to your installed CUDA version (nvcc --version). - For a deep learning-focused environment, use Python 3.9 with pip:
conda create -n deeplearning python=3.9 -y && pip install tensorflow[and-cuda] notebook