Create your own
Lesson illustration

Setting Up a Reproducible PyTorch Vision Environment on NVIDIA GPUs and Google Colab

Hello, and welcome to the first module of the course. We will build the practical PyTorch foundation needed for the modern vision systems later in the course: CNNs and transformers, detection and segmentation models, vision-language models, video pipelines, and deployment.

This first lesson establishes two working environments: a local NVIDIA GPU environment for repeatable development on your laptop, and a Google Colab GPU environment for workloads that exceed local VRAM or need a clean, shareable notebook runtime. The important principle is that “it ran once” is not enough: every run should record the software, hardware, and randomness settings that produced it.


1. What a reproducible vision environment contains

A vision experiment depends on more than a Python script. At minimum, capture:

  1. Code: Python files, notebooks, and configuration files under Git.
  2. Dependencies: Python and package versions, particularly the matched torch and torchvision build.
  3. Hardware/runtime facts: GPU model, NVIDIA driver version, CUDA runtime reported by PyTorch, and operating system.
  4. Randomness controls: seeds and deterministic-operation settings when you need comparable results.
  5. Data identity: eventually, dataset version, split files, and preprocessing configuration. We will construct those in later lessons.

A useful distinction:

  • Repeatability means rerunning on the same machine and environment produces the same result.
  • Reproducibility means another person, or your future self, has enough information to recreate the experiment closely.
  • Bitwise-identical results are not generally guaranteed across a GTX laptop and a Colab GPU. Floating-point implementations, GPU architectures, CUDA builds, and library versions can differ. The objective is controlled, explainable variation—not an unrealistic promise that every platform produces identical bits.

Create a project directory now:

cv-modern-vision/
├── notebooks/
├── src/
├── scripts/
├── environment/
├── data/
├── runs/
└── .gitignore

A sensible initial .gitignore is:

.venv/
__pycache__/
.ipynb_checkpoints/
data/
runs/
checkpoints/
*.pt
*.pth

Keep code, environment records, small configuration files, and dataset split manifests under version control. Do not casually commit raw datasets, model weights, API keys, or large generated outputs.


2. Local setup: isolated Python, NVIDIA driver, and the right PyTorch build

For your laptop, use a dedicated virtual environment rather than installing packages into the system Python or into a ROS environment. This prevents vision dependencies from quietly breaking other projects.

First, confirm that the NVIDIA driver can see your GPU. In a terminal, run:

nvidia-smi

You should see an NVIDIA GPU name, driver version, and VRAM. The driver is essential: it allows the operating system and PyTorch’s CUDA runtime to communicate with the GPU.

A common source of confusion is the local CUDA Toolkit. For ordinary PyTorch training and inference installed from official wheels, you usually do not need to install the full CUDA Toolkit separately. The PyTorch binary package supplies a compatible CUDA runtime. You do need a sufficiently current NVIDIA driver. A locally installed toolkit becomes relevant later if you compile custom CUDA extensions.

If nvidia-smi is absent or reports an error, update or reinstall the NVIDIA driver before proceeding. On Windows, the following short segment shows how to identify a laptop GPU through Task Manager and select its driver. The installation commands in the rest of this lesson should still come from the current official PyTorch selector, rather than from the video’s older package examples.

Install PyTorch for Windows GPU

“Install PyTorch for Windows GPU” by Jeff Heaton provides a quick driver-identification walkthrough for Windows systems.

If you are on Windows and do not know the precise GPU model, watch the driver setup. Focus on locating the dedicated NVIDIA GPU and selecting the matching notebook or desktop driver family.

Create the virtual environment

PyTorch currently requires a modern Python release; Python 3.10 or newer is appropriate. Python 3.11 is a good conservative choice for a new vision workspace.

Linux or WSL:

mkdir -p ~/cv-modern-vision
cd ~/cv-modern-vision

python3.11 -m venv .venv
source .venv/bin/activate

python -m pip install --upgrade pip

Windows PowerShell:

mkdir $HOME\cv-modern-vision
cd $HOME\cv-modern-vision

py -3.11 -m venv .venv
.\.venv\Scripts\Activate.ps1

python -m pip install --upgrade pip

When the environment is active, your prompt normally begins with (.venv). From this point onward, use python -m pip, not a bare pip, so the installer is unambiguously associated with the active interpreter.

Study the official selector before installing. It generates an installation command matched to your operating system and CUDA-capable configuration, and it changes as PyTorch releases change.

Get Started

Read PyTorch’s official installation guidance to select a current stable GPU-enabled wheel rather than relying on a copied command that may already be obsolete.

In the Start Locally section, use the interactive selector: choose Stable, your operating system, Pip, Python, and an NVIDIA CUDA compute platform. Read the installation guidance, then copy the command generated for your machine. Do not select CPU if you intend to use the laptop GPU. In the Installing on Linux or Installing on Windows section, review the With CUDA and Verification subsections for the platform-specific expectations.

Paste the command generated by the selector into your activated environment. It should install the matched torch, torchvision, and usually torchaudio packages.

Then add a small set of tools used throughout the course:

python -m pip install jupyterlab ipykernel numpy pillow matplotlib

Register this environment as a visible Jupyter kernel:

python -m ipykernel install --user --name cv-torch --display-name "Python (cv-torch)"

You can now launch Jupyter Lab from the activated environment:

jupyter lab

In Jupyter, select Python (cv-torch) as the notebook kernel. This avoids the frequent mistake of opening a notebook through a different Python installation that has a CPU-only or older PyTorch package.


3. Verify that PyTorch is actually using the GPU

A successful installation message does not prove that computation reaches the GPU. Create scripts/smoke_test.py:

import platform
import sys

import torch

print(f"Python: {sys.version.split()[0]}")
print(f"Platform: {platform.platform()}")
print(f"PyTorch: {torch.__version__}")
print(f"PyTorch CUDA runtime: {torch.version.cuda}")
print(f"CUDA available: {torch.cuda.is_available()}")

assert torch.cuda.is_available(), (
    "PyTorch cannot access CUDA. Check the NVIDIA driver, "
    "the active virtual environment, and the installed PyTorch build."
)

device = torch.device("cuda:0")
print(f"GPU: {torch.cuda.get_device_name(device)}")
print(f"cuDNN version: {torch.backends.cudnn.version()}")

x = torch.randn(1024, 1024, device=device)
y = x @ x
torch.cuda.synchronize()

assert y.is_cuda
print(f"Computation completed on: {y.device}")

Run it from the project root:

python scripts/smoke_test.py

A correct result includes:

  • CUDA available: True
  • the name of your NVIDIA GPU
  • Computation completed on: cuda:0

If torch.cuda.is_available() is False, diagnose in this order:

SymptomLikely causeAction
nvidia-smi failsDriver issue or GPU not exposed to the OSUpdate the NVIDIA driver; for WSL, verify GPU support is enabled in WSL.
nvidia-smi works, but CUDA is unavailable in PyTorchCPU build or wrong environment/kernelActivate .venv, check which python or where python, then reinstall using the official selector.
no kernel image is availableVery old GPU architecture unsupported by the installed buildCheck the GPU architecture; use Colab or a documented compatible legacy build if necessary.
Out-of-memory later in trainingModel, batch, or image size exceeds VRAMReduce batch size or image resolution; use Colab for larger experiments.

At this stage, do not benchmark your GPU from a single matrix multiplication. The real practical test is whether a later vision training loop fits in memory and has acceptable end-to-end throughput, including image loading and augmentation.


4. Google Colab: a portable GPU fallback, not a fixed machine

Colab provides a Jupyter-style notebook running on cloud infrastructure. It is useful for transformer fine-tuning, segmentation foundation models, and generative models that may not fit comfortably on a laptop GTX GPU.

The trade-off is that Colab is an ephemeral environment:

  • The assigned GPU can vary by session.
  • GPU access is subject to availability and account limits.
  • The runtime can disconnect or reset.
  • Packages and system images can change over time.
  • Files stored under /content disappear after a runtime reset.

The screenshot below shows the central action: selecting a GPU hardware accelerator in Colab. Available accelerators vary; the pictured T4 is an example, not a GPU type you can assume will always be assigned.

The Google Colab “Change runtime type” dialog, with Python 3 selected and the T4 GPU hardware accelerator selected. In practice, the available GPU choices depend on current Colab capacity and account access.

PyTorch using Google CoLab GPU to run your machine learning programs.

“PyTorch using Google CoLab GPU to run your machine learning programs” by ClarityCoders demonstrates the Colab runtime switch and a programmatic CUDA availability check.

Watch enabling the GPU. In the current Colab interface, open Runtime, choose Change runtime type, set Hardware accelerator to GPU, save, and then rerun cells from the top because changing runtimes resets execution state.

After enabling the GPU, make the first notebook cell a hardware assertion rather than silently falling back to CPU:

import platform
import sys
import torch

print(f"Python: {sys.version.split()[0]}")
print(f"Platform: {platform.platform()}")
print(f"PyTorch: {torch.__version__}")
print(f"PyTorch CUDA runtime: {torch.version.cuda}")

assert torch.cuda.is_available(), (
    "No CUDA GPU is attached. In Colab, use Runtime > Change runtime type "
    "and select GPU, then reconnect."
)

device = torch.device("cuda:0")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"GPU count: {torch.cuda.device_count()}")

In a second cell, inspect the NVIDIA-side information:

!nvidia-smi

Later, all models and batch tensors will need to be placed on the same selected device. The pattern is:

model = model.to(device)
images = images.to(device)
labels = labels.to(device)

For now, recognize the failure mode: a model on cuda:0 cannot operate directly on tensors left on the CPU. PyTorch reports this as a device mismatch. We will use this pattern deliberately in the next lesson.

Persist the important Colab artifacts

You may mount Google Drive to save notebooks, small configuration files, checkpoints, and environment records across runtime resets:

from google.colab import drive

drive.mount("/content/drive")

Then create a project folder, adjusting the path if you prefer a different Drive location:

from pathlib import Path

project_dir = Path("/content/drive/MyDrive/cv-modern-vision")
project_dir.mkdir(parents=True, exist_ok=True)

Only mount Drive in notebooks you trust: mounting grants the notebook access to files in that Drive account. Drive is convenient for persistence, but Git remains the better system of record for source code and configuration history.


5. Record the environment before you need to debug it

Once the local smoke test works, capture a baseline. Run these commands from the activated local environment:

python --version > environment/python-local.txt
python -m pip freeze --all > environment/pip-freeze-local.txt
nvidia-smi > environment/nvidia-smi-local.txt

For Colab, after mounting Drive, save a separate snapshot. Do not overwrite the local snapshot: they are intentionally different environments.

!python --version > /content/drive/MyDrive/cv-modern-vision/python-colab.txt
!python -m pip freeze --all > /content/drive/MyDrive/cv-modern-vision/pip-freeze-colab.txt
!nvidia-smi > /content/drive/MyDrive/cv-modern-vision/nvidia-smi-colab.txt

Also save PyTorch-specific metadata in a machine-readable form:

import json
import platform
import sys
from pathlib import Path

import torch

metadata = {
    "python": sys.version,
    "platform": platform.platform(),
    "torch": torch.__version__,
    "torch_cuda_runtime": torch.version.cuda,
    "cuda_available": torch.cuda.is_available(),
    "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
    "cudnn_version": torch.backends.cudnn.version(),
}

Path("environment").mkdir(exist_ok=True)
Path("environment/runtime-local.json").write_text(
    json.dumps(metadata, indent=2)
)

print(json.dumps(metadata, indent=2))

In Colab, change the final path to your mounted project_dir. This record will eventually accompany your training configurations and results.

A full pip freeze snapshot is an audit trail: it tells you exactly what was installed. It is not always portable across OSs and GPU types. Keep separate local and Colab snapshots, and later maintain a deliberately pinned project dependency file for packages that the project directly requires.


6. Seeds and deterministic execution

Training has multiple random sources: Python-level code, NumPy operations, PyTorch CPU operations, CUDA operations, randomized augmentation, shuffled data ordering, and sometimes nondeterministic GPU kernels.

PyTorch’s reproducibility documentation makes two key points:

  1. A seed gives a repeatable sequence of pseudo-random values within a controlled environment.
  2. Enforcing deterministic algorithms can reduce performance and may raise an error if a requested operation has no deterministic implementation.

Reproducibility — PyTorch 2.9 documentation

Read the official PyTorch notes to understand both the seed controls and the performance-versus-determinism trade-off on CUDA.

In Controlling sources of randomness, read the randomness controls, including the PyTorch, Python, NumPy, and CUDA convolution benchmarking subsections. Then read Avoiding nondeterministic algorithms, especially deterministic algorithm selection. Focus on why a fixed seed alone is insufficient when GPU algorithm selection can vary.

Use this function near the beginning of training scripts and notebooks:

import random

import numpy as np
import torch


def configure_reproducibility(seed: int = 2025, deterministic: bool = True) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if deterministic:
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        torch.use_deterministic_algorithms(True)

Call it before creating datasets, model weights, or data loaders:

configure_reproducibility(seed=2025, deterministic=True)

Use two explicit operating modes:

ModeUse caseConfiguration
Debug and comparisonInvestigating an error, checking whether a code change matters, producing a reported benchmarkFixed seed and deterministic algorithms enabled
Throughput explorationSearching for feasible batch sizes or profiling a stable pipelineFixed seed retained; determinism may be relaxed consciously and recorded

A deterministic configuration is not “better” in every situation. It trades some speed and available kernels for a clearer causal story: if a result changes, you have fewer uncontrolled explanations.

For stricter CUDA reproducibility, set process-level variables before starting Python or Jupyter. This is particularly useful when a CUDA operation reports a cuBLAS determinism warning or error.

Linux/WSL:

export PYTHONHASHSEED=2025
export CUBLAS_WORKSPACE_CONFIG=:4096:8
jupyter lab

Windows PowerShell:

$env:PYTHONHASHSEED = "2025"
$env:CUBLAS_WORKSPACE_CONFIG = ":4096:8"
jupyter lab

These values must be recorded alongside results. Changing the seed is legitimate for robustness studies, but it should be an intentional experimental variable, not an accidental change between notebook runs.


Completion checklist

Before moving on, verify that you can check every item:

  • A local .venv exists and contains GPU-enabled PyTorch and torchvision.
  • nvidia-smi works locally and scripts/smoke_test.py reports CUDA available: True.
  • Jupyter exposes the Python (cv-torch) kernel.
  • A Colab notebook can request a GPU and asserts that CUDA is available.
  • You have saved separate local and Colab package, Python, and GPU snapshots.
  • Your project contains a reusable configure_reproducibility() function.
  • You understand that strict deterministic execution is a controlled experimental mode, not a universal guarantee across different machines.

You now have two viable execution targets and a baseline record for both. The essential habit is simple: assert the device, isolate the environment, capture versions, and control randomness before you trust a result.

Next, we will work directly with PyTorch tensors: batched image-shaped data, GPU placement, vectorized operations, and automatic differentiation.

Can't find a good explanation? Sign up and we'll make it for you

Sign up