Skip to content

Instantly share code, notes, and snippets.

@ShardulJunagade

ShardulJunagade/PyTorch GPU Setup.md

Created February 8, 2025 14:54
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save ShardulJunagade/0de582ec00bbf45a275fdcd00eb00592 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save ShardulJunagade/0de582ec00bbf45a275fdcd00eb00592 to your computer and use it in GitHub Desktop.
Download ZIP
Step-by-step guide to installing PyTorch with NVIDIA GPU support using venv, Conda, or Docker.
Raw
PyTorch GPU Setup.md

Setting Up PyTorch with NVIDIA GPU (CUDA and cuDNN)

This guide provides three different methods to install PyTorch with GPU acceleration using CUDA and cuDNN. Choose the method that best suits your requirements and system configuration.

  1. Using Virtual Environment (venv) in Python
  2. Using Conda Environment in Anaconda
    • Using pip
    • Using conda
  3. Using Docker Container

Check the official PyTorch website for the latest installation instructions.

Method 1: Using Virtual Environment (venv) in Python

This method sets up PyTorch with CUDA using a Python virtual environment.

Steps

  1. Install the latest version of Python from the official website.

  2. Create a virtual environment:

    python -m venv mlenv

  3. Activate the virtual environment:

    • Windows: 
      mlenv\Scripts\activate
    • Linux/macOS: 
      source mlenv/bin/activate

  4. Install PyTorch with CUDA support:

    First, check your GPU and CUDA version using nvidia-smi. If you don't have CUDA installed, download CUDA Toolkit and cuDNN from the NVIDIA website.

    Then, install PyTorch with the appropriate CUDA version:

    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

  5. Verify the installation:

    python -c "import torch; print(torch.cuda.is_available())"

    If the output is True, PyTorch has been successfully installed with GPU support.

  6. Deactivate the virtual environment:

    deactivate

  7. Remove the virtual environment (optional):

    • Windows: 
      rmdir /s /q mlenv
    • Linux/macOS: 
      rm -rf mlenv

Method 2: Using Conda Environment in Anaconda

This method sets up PyTorch with CUDA using a Conda environment.

Steps

  1. Install Anaconda or Miniconda.

  2. Add Conda to the system PATH:

    • Windows:
      • User Installation Add the following directories to your user's environment variables: 
        C:\Users\YourUsername\anaconda3
C:\Users\YourUsername\anaconda3\Scripts
C:\Users\YourUsername\anaconda3\Library\bin
      • System Installation Add the following directories to your system's environment variables: 
        C:\ProgramData\Anaconda3
C:\ProgramData\Anaconda3\Scripts
C:\ProgramData\Anaconda3\Library\bin
    • Linux/macOS: Add the following line to your shell configuration file (e.g., ~/.bashrc, ~/.bash_profile, ~/.zshrc): 
      export PATH="/path/to/anaconda3/bin:$PATH"
      Then, run: 
      source ~/.bashrc    # or ~/.bash_profile, ~/.zshrc

  3. Create a Conda environment:

    conda create -n mlenv python=3.12 -y
conda activate mlenv

  4. Install PyTorch with CUDA support:

    First, check your GPU and CUDA version using nvidia-smi. Then, install PyTorch with the appropriate CUDA version:

    • using pip:

      If you don't have CUDA installed, download CUDA Toolkit and cuDNN from the NVIDIA website.

      pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

    • using conda:

      This method ensures CUDA and cuDNN are installed via Conda

      conda install cudatoolkit -c anaconda -y
conda install pytorch-cuda=12.4 -c pytorch -c nvidia -y
conda install pytorch torchvision torchaudio -c pytorch -c nvidia -y

  5. Verify the installation:

    python -c "import torch; print(torch.cuda.is_available())"

    If the output is True, PyTorch has been successfully installed with GPU support.

  6. Deactivate the Conda environment:

    conda deactivate mlenv

  7. Remove the Conda environment (optional):

    conda env remove -n mlenv

Method 3: Using Docker Container

This method sets up PyTorch with CUDA using a Docker container.

Steps

  1. Install Docker Desktop on your system.

  2. Enable GPU support in Docker Desktop:

    • Open Docker Desktop.
    • Go to Settings > Resources > Advanced.
    • Enable Use the GPU and Enable experimental features.

  3. Ensure that Docker has access to the GPU:

  4. Pull the PyTorch Docker image with GPU support:

    docker pull pytorch/pytorch:latest

    This command downloads the latest PyTorch image.

  5. Run the Docker container:

    docker run --gpus all -it --rm --shm-size=8g pytorch/pytorch:latest

    This command starts a Docker container with GPU support and an interactive terminal.

    • --gpus all -> Enables GPU support.
    • -it -> Interactive mode.
    • --rm -> Automatically removes the container when it is stopped.
    • --shm-size=8g -> Sets the shared memory size to 8 GB to avoid shared memory issues.

    Alternatively, if you want a persistent container with access to local files, run:

    docker run --gpus all -it --name pytorch_container --shm-size=8g -v $(pwd):/workspace pytorch/pytorch:latest

    This command mounts the current directory to the /workspace directory in the container.

    • --name pytorch_container -> Names the container.
    • -v $(pwd):/workspace -> Mounts the current directory to the /workspace directory in the container.

  6. Verify the installation:

    python -c "import torch; print(torch.cuda.is_available())"

    If the output is True, PyTorch has been successfully installed with GPU support.

  7. Exit the Docker container:

    exit

    This command exits the Docker container.

  8. Remove the Docker container (optional):

    docker rmi pytorch/pytorch:latest

    This command removes the PyTorch Docker image.

Additional Docker Commands

  • Restart a stopped container: 
    docker start -ai pytorch_container
  • Remove a stopped container: 
    docker rm pytorch_container

Verifying Installation

To verify that PyTorch has been successfully installed with GPU support, run the following Python code:

import torch
print("CUDA available:", torch.cuda.is_available())     # Returns True if CUDA is available
if torch.cuda.is_available():
    current_device = torch.cuda.current_device()
    print("Current device ID:", current_device)
    print("Current device name:", torch.cuda.get_device_name(current_device))
    print("Device memory address:", torch.cuda.device(current_device))
    print("Total number of GPUs:", torch.cuda.device_count())

References

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment