Stop 01 · Programming & CS FoundationsTopics · 4/12

OOP in Python

Part of the AI Engineer Roadmap

In short

Open the source of scikit-learn or PyTorch and you'll find classes everywhere. Model objects with fit() and predict() only make sense once you understand how classes and objects work.

Classes, objects, inheritance and dunder methods — most ML libraries (scikit-learn estimators, PyTorch modules) are built around this object model, so reading their source requires it.

The purple blob pulls pista, mango and rose kulfis out of one metal kulfi mould.

Think of it like this

One kulfi mould, many kulfis

The kulfi wala has one metal mould. Every kulfi that comes out has the same cone shape, but one is pista, one is mango, one is rose. The class is the mould. Each object is a kulfi with its own flavour, made from the same design.

Watch it happen

A class family and its objects

Model
LinearModel
TreeModel
price_model
rent_model
churn_model

1/4Model is the parent class. It defines what every model must have, like fit() and predict().

The big ideas

What you need to know

A class bundles data and actions

Attributes hold data (self.weights). Methods are functions that act on it (self.fit()). They travel together in one object.

__init__ sets up each object

__init__ runs once when you create the object. self is the object itself, so self.lr = lr stores the value on it.

Inheritance reuses code

class TreeModel(Model) gets everything Model has. Override only the methods that need to behave differently.

Dunder methods add superpowers

__len__, __repr__ and __getitem__ let your objects work with len(), print() and square brackets. PyTorch datasets rely on these.

See it in code

Walk through the code

python
1class MeanPredictor:2    """Predicts the average of whatever it was trained on."""3 4    def __init__(self, name="baseline"):5        self.name = name6        self.mean_ = None7 8    def fit(self, y):9        self.mean_ = sum(y) / len(y)10        return self11 12    def predict(self, n):13        return [self.mean_] * n14 15    def __repr__(self):16        return f"MeanPredictor(name={self.name!r}, mean={self.mean_})"

1/4The class is the mould. Nothing has been created yet.

This follows the scikit-learn style: fit() learns and returns self, predict() uses what was learned, and the trailing underscore in mean_ marks a learned value.

Save yourself the pain

Common mistakes

Oops

Forgetting self in a method, like def fit(y):.

Do this instead

Every normal method takes self first. Python passes the object in for you when you call model.fit(y).

Oops

Making a class for everything, even a two-line helper.

Do this instead

Use a plain function when there's no state to keep. Reach for a class when data and behaviour belong together.

Your mission

Write your own baseline model

~25 min

Quick check

You create a = MeanPredictor() and b = MeanPredictor(), then call a.fit([1, 2, 3]). What is b.mean_?

Go deeper

  1. 1Build a small class hierarchy (e.g. Shape → Circle, Square) to practice inheritance.
  2. 2Read the source of a simple scikit-learn estimator to see OOP conventions in practice.
  3. 3Learn `__init__`, `__call__` and `__repr__` — the dunder methods you'll see constantly in PyTorch code.

More topics at this stop

Stuck on this topic? Ask an Insider

Get 1:1 guidance from people who've walked this exact path — on the InsideEdge app.

Download
InsideEdge