Search⌘ K
AI Features

Going Classy

Explore how to create and manage PyTorch classes that organize model training into configurable components. Learn to define constructors with essential and optional arguments, set placeholders for data loaders, and track training variables to improve code modularity and flexibility. This lesson helps you build classes that simplify changing parameters like epochs, optimizers, and devices without modifying the core training loop.

Building classes for model training

So far, the %%writefile magic has helped us to organize the code into three distinct parts: data preparation, model configuration, and model training. At the end of the chapter, Rethinking the Training Loop, though, we bumped into some of its limitations, e.g., being unable to choose a different number of epochs without editing the model training code.

Clearly, this situation is not ideal. We need to do better. We need to go classy; that is, we need to build a class to handle the model training part.

The class

Let us start defining our class with a rather unoriginal name: StepByStep. Since we are starting it from scratch, either we do not specify a parent class or we inherit it from the fundamental object class. I prefer the latter, so our class definition looks like this:

Python 3.5
# A completely empty (and useless) class
class StepByStep(object):
pass

Boring, right? Let us make it more interesting by building it further.

The constructor

...

“From where do we start building a class?” ...