Setting up a Python ML environment: a step-by-step guide

Introduction

Before you train your first model or tune hyperparameters, you need a solid foundation. Setting up a Python ML environment: a step-by-step guide is about establishing the tools, isolation, and workflows that let you move fast without breaking dependencies or losing reproducibility. We've seen too many projects stall because a junior engineer installed libraries globally, or a senior engineer spent an afternoon debugging why TensorFlow wouldn't load on the GPU.

The good news is that the setup itself takes less than ten minutes. The hard part—and the part that matters—is doing it right from day one.

Core libraries: the working minimum

Most ML work in Python starts with four or five libraries that have become the industry standard. NumPy handles numerical operations and arrays. Pandas gives you DataFrames—the mental model most of us use for tabular data. Scikit-learn provides classic machine learning algorithms and preprocessing tools. Matplotlib (or Seaborn for prettier plots) handles visualization. Jupyter notebooks let you iterate and explore without restarting a script every three seconds.

If you're building a recommendation engine or a fraud detection model, these five tools will take you 80% of the way there. They're battle-tested, well-documented, and widely supported in production environments. More importantly, they work well together.

NumPy underpins everything. Pandas builds on NumPy. Scikit-learn uses both. Jupyter ties it all together in an interactive environment. That's the mental model: layers of abstraction, each one solving a specific problem.

Here's what a typical workflow looks like: you load a CSV with pandas, explore the data in a Jupyter notebook, clean it, engineer features using NumPy operations, train a scikit-learn classifier, and visualize results with matplotlib. No external APIs, no spinning up containers, no waiting for cloud resources to provision. Just pure computation on your laptop.

Isolate dependencies: one virtual environment per project

This is where most setups go wrong. Installing packages globally sounds convenient until you have ten projects, each depending on different versions of the same library. Project A needs scikit-learn 0.24, but project B requires 1.0 and they're incompatible. Now you're stuck.

Virtual environments solve this by creating an isolated Python installation per project. You switch between projects, and each one has its own dependency tree. Lock files (requirements.txt, or better yet, a pyproject.toml with version pinning) ensure that when your teammate clones the repo next week, they get the exact same versions you have today.

We recommend uv for this. It's a fast Python package installer that handles both installation and virtual environment management in one tool. Here's the workflow:

uv init ml-lab

This creates a new project directory called ml-lab with a minimal pyproject.toml. Inside that directory, run:

uv add numpy pandas scikit-learn matplotlib jupyter

Behind the scenes, uv creates a virtual environment and installs each package into it. It also generates a lock file (uv.lock) that pins exact versions. When you check that lock file into Git, your teammates and CI/CD pipeline will install the same dependencies.

To activate the virtual environment on your machine, uv provides a simple wrapper:

uv run python script.py

Or for interactive work:

uv run jupyter notebook

No need to remember activation commands. uv handles it. If you're already comfortable with venv and pip, the old way still works—but uv is faster and more reliable, especially when dealing with complex dependency graphs.

GPU when needed: PyTorch with CUDA for deep learning

Not every ML project needs a GPU. If you're training a random forest on a million rows of tabular data, your laptop's CPU is fine. Scikit-learn, pandas, and NumPy all run happily without GPUs.

But if you're building neural networks, language models, or computer vision systems, you want GPU acceleration. This is where PyTorch comes in. And here's the key: only add PyTorch when you need it. Don't bloat your environment with deep learning libraries if you're doing classic supervised learning.

When you do need it, the installation depends on your hardware. NVIDIA GPUs use CUDA. Apple Silicon uses Metal Performance Shaders (built into PyTorch). AMD GPUs use ROCm. Installing the wrong variant wastes hours of debugging.

PyTorch's official website has a configurator. You select your OS, package manager, language, and compute platform, and it gives you the exact command. For an Ubuntu machine with an NVIDIA GPU, it might look like:

uv add torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

The cu118 flag specifies CUDA 11.8. For CPU-only (useful for testing on a machine without a GPU), you'd use a different URL. For Mac with Apple Silicon, PyTorch detects Metal automatically—no extra flags needed.

After installation, verify that PyTorch can see your GPU:

import torch
print(torch.cuda.is_available())  # Should print True if GPU is detected
print(torch.cuda.get_device_name(0))  # Should print your GPU's name

If that returns False, your installation didn't work. Check that you used the right CUDA version for your system (nvidia-smi will tell you which CUDA version your GPU supports).

Notebooks vs scripts: exploration and production code

Jupyter notebooks are incredible for exploration. You run a cell, see the output, tweak the code, run it again. You can write prose alongside code, create visualizations inline, and keep your thinking process visible. Early-stage work belongs in notebooks.

But notebooks have a fatal flaw: they're hard to test, version control, and deploy. A cell can run out of order. Variables leak between cells. Refactoring is error-prone. Notebooks don't scale to production.

The pattern we follow is: explore in a notebook until you understand the problem, then move reusable code into Python modules. A module is just a .py file with functions and classes. It's testable, version-controllable, and importable from anywhere.

Here's a concrete example. You start in exploration.ipynb, loading data and training a classifier:

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.train_test_split import train_test_split

df = pd.read_csv('data.csv')
X = df.drop('target', axis=1)
y = df['target']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
clf = RandomForestClassifier(n_estimators=100, max_depth=10)
clf.fit(X_train, y_train)
score = clf.score(X_test, y_test)
print(f"Accuracy: {score}")

Once you've validated this approach, extract the reusable parts into a module. Create ml_lab/model.py:

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

def train_classifier(data_path, test_size=0.2, n_estimators=100, max_depth=10):
    df = pd.read_csv(data_path)
    X = df.drop('target', axis=1)
    y = df['target']
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size)
    clf = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth)
    clf.fit(X_train, y_train)
    
    return clf, X_test, y_test

Now your notebook becomes a consumer of that module:

from ml_lab.model import train_classifier

clf, X_test, y_test = train_classifier('data.csv')
score = clf.score(X_test, y_test)
print(f"Accuracy: {score}")

Later, when you need to train this model in a batch job or API endpoint, you import the same function. No code duplication. No notebook magic. This is how teams scale ML work from prototype to production.

Common mistakes

• Installing globally. "pip install numpy" in your system Python instead of in a virtual environment. Weeks later, your projects break because a dependency changed. Always use a virtual environment.
• Ignoring lock files. Committing only pyproject.toml without the lock file means your teammates might install different versions. Lock files (uv.lock or requirements.txt) must go in version control.
• GPU misconfigurations. Installing PyTorch CPU when you have an NVIDIA GPU, or vice versa. Always verify with torch.cuda.is_available() immediately after installation.
• Keeping everything in notebooks. A 500-line notebook is a code smell. When you have reusable logic (data loading, preprocessing, model training), move it to modules early.
• Bloated environments. Adding every ML library under the sun "just in case." Start lean. Add dependencies when you need them. Smaller environments are faster to install and easier to debug.

Key takeaways

• NumPy, pandas, scikit-learn, matplotlib, and Jupyter form the foundation for most ML work. Master these five libraries before reaching for anything else.
• Use a virtual environment per project and commit a lock file to version control. uv is a modern, fast alternative to venv and pip.
• Add PyTorch (or TensorFlow) only when you need deep learning. Verify GPU support immediately after installation with torch.cuda.is_available().
• Start exploration in Jupyter notebooks, but move reusable code into Python modules. Modules are testable, version-controllable, and deployable.
• Repeatability is the foundation of good ML work. If your setup takes a colleague five minutes to reproduce, you've done it right.