We've got a training class - let's write some code to use this training class and train the Pix2Pix architecture. Here are the steps needed to run the training class and change the settings:

1. Create a run.py file in the src folder that has the following code inside of it:

#!/usr/bin/env python3
from train import Trainer

# Command Line Argument Method
HEIGHT = 256
WIDTH = 256
CHANNELS = 3
EPOCHS = 100
BATCH = 1
CHECKPOINT = 50
TRAIN_PATH = "/data/cityscapes/cityscapes/train/"
TEST_PATH = "/data/cityscapes/cityscapes/val/"

trainer = Trainer(height=HEIGHT,width=WIDTH, channels=CHANNELS,epochs =EPOCHS,
                 batch=BATCH,
                 checkpoint=CHECKPOINT,
                 train_data_path=TRAIN_PATH,
                 test_data_path=TEST_PATH)
trainer.train()

The training class has the following inputs:

• • HEIGHT: Height of the input images
  • WIDTH: Width of the input images
  • CHANNELS: Number of Channels to the input image
  • EPOCHS: Number of times we will train the model with the dataset
  • BATCH: How many images per forward pass through the network
  • CHECKPOINT: How often do we want to check the model outputs
  • TRAIN_PATH: The path to the training data
  • TEST_PATH: The path to the test data

2. Next, create a file called run.sh script at the root of the code directory and add the following text into the file:

#/bin/bash

# Training Step
xhost +
docker run -it 
   --runtime=nvidia 
   --rm 
   -e DISPLAY=$DISPLAY 
   -v /tmp/.X11-unix:/tmp/.X11-unix 
   -v $HOME/Chapter5/out:/out 
   -v $HOME/Chapter5/src:/src 
   ch5 python3 /src/run.py

3. At the root of this chapter's repository, run the following script with the sudo command (make sure the shell script is executable):

sudo ./run.sh

Your model should be training now - congratulations! You have successfully implemented Pix2Pix from scratch.