Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Training [3/4] | by Johannes Schmidt | Towards Data Science

In the previous chapters we created our dataset and built the U-Net model. Now it is time to start training. For that we will write our own training loop within a simple Trainer class and save it in trainer.py. The Jupyter notebook can be found here. The idea is that we can instantiate a Trainer object with parameters such as the model, a criterion etc. and then call it’s class method run_trainer() to start training. This method will output the accumulated training loss, the validation loss, and the learning rate that was used for training. Here is the code:
In order to create a trainer object the following parameters are required:
model: e.g. the U-Net
device: CPU or GPU
criterion: loss function (e.g. CrossEntropyLoss, DiceCoefficientLoss)
optimizer: e.g. SGD
training_DataLoader: a training dataloader
validation_DataLoader: a validation dataloader
lr_scheduler: a learning rate scheduler (optional)
epochs: The number of epochs we want to train
epoch: The epoch number from where training should start
Training can then be started with the class method run_trainer(). Since training is usually performed with a training and a validation phase, _train() and _validate() are two functions that are run once for every epoch we train with run_trainer() (line 33–53). If we have a lr_scheduler, we also perform a step with the lr_scheduler. To visualize the progress of training, I included a progress bar with the library tqdm. Now let’s take a closer look on what happens when calling _train() and _validate() . If you are familiar with using PyTorch for network training, there is probably nothing new here.
In _train() we basically just iterate over our training dataloader and send our batches through the network in train mode (line 56–64). We then use this output together with our target to compute the loss with the loss function for the current batch (line 65). The computed loss is then appended in a temporary list (line 66–67). Based on the computed gradients, we perform a backward pass and a step with our optimizer to update the model’s parameters (line 68–69). At the end we update our progress bar for the training phase to show the current loss (line 71). The function outputs the mean of the temporary loss list and the learning rate that was used.
In _validate() , similar to _train() , we iterate over our validation dataloader, send our batches through the network in validation mode and compute the loss. This time, without computing the gradients and without performing a backward pass (line 78–97).
Let’s create our Carvana data generators once again, but this time run the code within a Jupyter notebook.
Please note that I resize the images to 128x128x3 using skimage.transform.resize() to speed up training. This will generate batches of images that look like this:
from visual import DatasetViewerdataset_viewer_training = DatasetViewer(dataset_train)dataset_viewer_training.napari()
I can then instantiate the Trainer object and start training:
Training will look something like this:
Although training was performed on a NVIDIA 1070, it took 1:19 min to train 2 epochs with only 96 images (size 128x128x3) for each epoch. Why is that? The reason why this is so painfully slow, is because every time we generate a batch we read the data in full resolution (1918x1280x3) and resize it. And we do this for every epoch! Therefore, it would make more sense to either store the data in a lower resolution and then to pick the data up, or store the data in cache and access it when it’s needed. Or both. Let’s slightly change our custom SegmentationDataSet1 class (create a new class called SegmentationDataSet2):
Here we added the argument use_cache and pre_transform . We basically just iterate over our input and target list and store the images in a list when we instantiate our dataset. When __getitem__ is called, an image-target pair from this list is returned. I added the pre_transform argument because I don’t want to change the original files. Instead, I want the images to be picked up, resized and stored in memory. Again, I included a progress bar to visualize the caching. Let’s try it out. The changes in code are the following:
And it looks something like this:
The first progress bar represents the training dataloader and the second the validation dataloader. Let’s train again for 2 epochs and see how long it’ll take.
Training took about 2 seconds only! That’s much better. But there is one part we can still improve. Creating the dataset that reads images and stores them in memory takes a bit of time. When you look at the code and the CPU usage, you’ll notice that only one core is used. Let’s change it in a way, so that all cores are used. Here I use the multiprocessing library:
Please note that there is no progressbar in this dataset class!
Before we perform training, let’s also make a quick detour and talk about the learning rate.
The learning rate is one of the most important hyperparameters in neural network training. Choosing proper learning rates throughout the learning procedure is difficult as a small learning rate leads to slow convergence while a high learning rate can cause divergence. Also, frequent parameter updates with high variance in SGD can cause fluctuations, which makes finding the (local) minimum for SGD even more difficult. To identify an optimal learning rate, we can test different learning rates empirically with a learning rate range test. Inspired by the best practices I picked up from the fast.ai course, I recommend using a learning rate finder before starting the actual training. Sylvain Gugger from fast.ai wrote a really good summary about this problem. The code that I will show you is based on Tanjid Hasan Tonmoy’s pytorch-lr-finder, which is an implementation of the learning rate range test from Leslie Smith. I only slightly modified the code and included a progressbar (yes, I like them).
Let’s perform such a learning rate range test. Since our dataset is rather small (96 images), we’ll perform some extra steps (1000). The upper progressbar displays the number of epochs and the lower progressbar shows the number of steps we perform on the current epoch.
Let’s plot the results of the test:
0.01 seems to be a good learning rate. We’ll take it. Let’s train for 100 epochs...
...and visualize the training and validation loss. For that I will use matplotlib and write a function that I can add to the visual.py file.
Let’s see what the function plot_training() will output when we pass in our losses and the learning rate.
Training looks good!
We can then save our model with PyTorch:
# save the modelmodel_name =  'carvana_model.pt'torch.save(model.state_dict(), pathlib.Path.cwd() / model_name)
In this part, we performed training with a sample of the Carvana dataset by creating a simple training loop. The progress of this training loop can be visualized with a progressbar and the result of training can be plotted with matplotlib. We noticed that training was painfully slow because our data was picked up very slowly by our custom data generator. Because of that, we changed it in a way so that data is only read once and then picked up from memory when needed. We also made use of multiprocessing for that case. Additionally, we added a learning rate range finder, to determine an optimal learning rate which we then used for model training.
In the next chapter, we’ll let the model predict the segmentation maps of unseen image data (inference).