Getting hands-on with AI
The first code we’ll examine here is the actual model itself. This code defines the neural network and how it’s structured, as well as how it’s called. All of this is stored in the lib/models.py
library file.
First, we load any libraries we’re using:
import torch from torch import nn
In this case, we only import PyTorch and its nn
submodule. This is because we only include the model code in this file and any other libraries will be called in the file that uses those functions.
Defining our upscaler
One of the most important parts of our model is the upscaling layers. Because this is used multiple times in both the encoder and decoder, we’ve broken it out into its own definition, and we’ll cover that here:
- First, we define our class:
class Upscale(nn.Module): """ Upscale block to double the width/height from depth. """ def __init__(self, size...