How To Structure Your PyTorch Project | by Branislav Holländer | Towards Data Science

Ever since I started to train deep neural networks, I was wondering what would be the structure for all my Python code. Ideally, a good structure should support extensive experimenting with the model, allow for implementing various different models in a single compact framework and be easily understood by everybody reading the code. You have to be able to use data from different data sources by coding and reusing various data loaders. Additionally, it would be nice if the models supported combining multiple networks in one model (such as is the case for GANs or the original R-CNN). The framework should also have enough flexibility to allow for complex visualizations (it is one of my core beliefs in data science that visualization makes everything MUCH easier, especially in the case of computer vision tasks).
The detailed implementation of a deep learning framework will of course be dependent on the underlying library you are using, whether it be TensorFlow, PyTorch or CNTK. In this post I will present my approach based on PyTorch. However, I think that the general structure applies equally to whatever libraries you are using. You can find the whole repository at https://github.com/branislav1991/PyTorchProjectFramework.
On the image above (taken from VS code, my Python editor of choice), you can see the general folder structure that I created for my framework. The framework consists of some startup scripts (train.py, validate.py, hyperopt.py) as well as the libraries hiding inside the folders. The datasets folder contains classes and methods for loading various types of data for training. The losses folder may contain additional loss functions or validation metrics. If you do not require any custom loss functions for your project, you probably won’t need this folder. The models folder is the most important: it contains the actual models. The optimizers folder includes code for custom optimizers. As with the losses folder, if you do not have any custom optimizers, you may as well omit this folder. Finally, the utils folder contains various utilities used all over the framework, most notably the visualizer. You will also notice the config_segmentation.json file in the root folder of the project. This file contains all the configuration options required for training.
As you may have guessed, training is launched by calling the train.py script. This script is called with the appropriate configuration file as a command line argument. It takes care of all the high-level training stuff such as loading the training and validation datasets and the model, setting up the visualization, running the training loop and exporting the trained model at the end.
Validation is used similarly by calling the appropriate scripts and passing the configuration files as arguments.
On the image above you can see the structure of the dataset folder. It includes the __init__.py module which includes some necessary functions to find and create the right dataset as well as a custom data loader which forwards the data to the training pipeline (for more information on this, please have a look at the PyTorch API documentation). The base_dataset.py, as the name suggests, defines the abstract base class for each dataset you define in the framework.
For every custom dataset you define, you will have to implement the __getitem__ and __len__ methods so that PyTorch may iterate over it. You won’t have to deal with the DataLoader anymore since that is defined in datasets/__init__.py already. You may also define custom callbacks for the dataset to be called before and after every epoch. This could be useful if you want to use some warmup method that feeds different data to the model during the first few epochs and switches to a more complex dataset afterwards.
To instantiate the dataset, the train.py script calls the following code:
print(‘Initializing dataset...’)train_dataset =   create_dataset(configuration[‘train_dataset_params’])train_dataset_size = len(train_dataset)print(‘The number of training samples = {0}’.format(train_dataset_size))
This calls the create_dataset function which looks at the configuration file and chooses the correct dataset based on its name. It is important to follow the convention <datasetname>_dataset.py when naming the dataset since that is how the script is able to find your dataset based on the string in the configuration file. Finally, the script above calls the len() function on the dataset to inform you about its size.
Models in the framework work in the same way as the datasets: the __init__.py module includes functions to find and create the right model according to its module name and the string defined in the configuration file. The model class itself inherits from the abstract BaseModel class and has to implement two methods:
forward(self) to run the forward prediction, and
optimize_parameters(self) to modify the weights of the network(s) after a training pass.
All other methods may either be overridden or you may use the default BaseClass implementation. Functions you may want to override include pre_epoch_callback and post_epoch_callback (called before and after each epoch) or test (called during validation).
In order to use the framework correctly, it is important to know how to use the networks, the optimizers and the losses in the model. Since there may be multiple networks in a model using different optimizers as well as multiple different losses (for instance, you could want to display the bounding box classification and regression losses for a semantic localization model), the interface is a bit more involved. Specifically, you provide the names of your losses and your networks together with the optimizers for the BaseModel class to know how to train your model. In the provided code I included an example of a 2D segmentation model together with an example dataset for you to see how the framework is supposed to be used.
Take a look at the __init__() function of the provided 2D segmentation model:
class Segmentation2DModel(BaseModel):   def __init__(self, configuration):      super().__init__(configuration)      self.loss_names = [‘segmentation’]      self.network_names = [‘unet’]      self.netunet = UNet(1, 2)      self.netunet = self.netunet.to(self.device)      if self.is_train: # only defined during training time         self.criterion_loss = torch.nn.CrossEntropyLoss()         self.optimizer = torch.optim.Adam(self.netunet.parameters(), lr=configuration[‘lr’])         self.optimizers = [self.optimizer]
This is what is happening here: first, we read the model configuration. Then, we define the ‘segmentation’ loss and put it in the self.loss_names list. The name of the loss is important since we use the variable self.loss_segmentation for the loss. By knowing the name, the BaseModel can look up the loss and print it in the console or visualize it (more about visualization in the next section). Similarly, we define the name of the network. This makes sure that BaseModel knows how to train the model without us having to explicitly define it. Next, we initialize the network (in this case a U-Net) and move it to the GPU. If we are in training mode, we also define the loss criterion and instantiate the optimizer (in this case Adam). Lastly, we put the optimizer into the self.optimizers list. This list is again used in the BaseModel class to update the learning rate or to resume training from a given checkpoint.
Let us also take a look at the forward() and optimize_parameters() functions:
def forward(self):   self.output = self.netunet(self.input)def backward(self):   self.loss_segmentation = self.criterion_loss(self.output, self.label)def optimize_parameters(self):   self.loss_segmentation.backward() # calculate gradients   self.optimizer.step()   self.optimizer.zero_grad()
As you can see, this is standard PyTorch code: its only responsibility is to call forward() on the network itself, to step the optimizer after the gradients have been calculated and to zero them again. It should be easy to implement this for your own model.
Visualization can be found in the Visualizer class. This class is responsible for printing out loss information to the terminal as well as visualizing various results using the visdom library. It is initialized at the beginning of the training script (which loads up the visdom server). The training script also calls its plot_current_losses() and print_current_losses() functions to visualize and write out the training loss. It also contains functions like plot_current_validation_metrics(), plot_roc_curve() and show_validation_images() which are not called automatically but may be called from the model in the post_epoch_callback() to do some useful visualization upon validation. I tried to keep the visualizer fairly general. Certainly you can expand the functionality of the visualizer yourself to make it more useful for you.
I presented an approach to writing a general deep learning framework that can be used in all areas of deep learning. By using this structure, you will obtain a clear and flexible codebase for further development. Of course there are many alternative ways how to approach the problem. Let me know in the comments if you have other suggestions!