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:
- Code: Python files, notebooks, and configuration files under Git.
- Dependencies: Python and package versions, particularly the matched
torchandtorchvisionbuild. - Hardware/runtime facts: GPU model, NVIDIA driver version, CUDA runtime reported by PyTorch, and operating system.
- Randomness controls: seeds and deterministic-operation settings when you need comparable results.
- 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.
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:
| Symptom | Likely cause | Action |
|---|---|---|
nvidia-smi fails | Driver issue or GPU not exposed to the OS | Update the NVIDIA driver; for WSL, verify GPU support is enabled in WSL. |
nvidia-smi works, but CUDA is unavailable in PyTorch | CPU build or wrong environment/kernel | Activate .venv, check which python or where python, then reinstall using the official selector. |
no kernel image is available | Very old GPU architecture unsupported by the installed build | Check the GPU architecture; use Colab or a documented compatible legacy build if necessary. |
| Out-of-memory later in training | Model, batch, or image size exceeds VRAM | Reduce 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
/contentdisappear 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.

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:
- A seed gives a repeatable sequence of pseudo-random values within a controlled environment.
- 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:
| Mode | Use case | Configuration |
|---|---|---|
| Debug and comparison | Investigating an error, checking whether a code change matters, producing a reported benchmark | Fixed seed and deterministic algorithms enabled |
| Throughput exploration | Searching for feasible batch sizes or profiling a stable pipeline | Fixed 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
.venvexists and contains GPU-enabled PyTorch andtorchvision. -
nvidia-smiworks locally andscripts/smoke_test.pyreportsCUDA 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