"""Trains the FRRN A architecture on the CityScapes dataset."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import collections
import functools
import pickle
import lasagne
import logging as pylogging
from dltools import hybrid_training, utility, architectures, logging, optimizer, hooks, losses
import theano
import theano.tensor as T
import numpy as np
# Try to import pychianti and if it doesn't work, import the python adapter.
# This allows us to use the code without installing the C++ library.
try:
import pychianti
except:
from dltools import pychianti_adapter as pychianti
NUM_CLASSES = 19
IMAGE_CHANNELS = 3
IMAGE_ROWS = 1024
IMAGE_COLS = 2048
BASE_CHANNELS = 48
FR_CHANNELS = 32
MULTIPLIER = 2
BOOTSTRAP_MULTIPLIER = 64
REDUCE_LR_INTERVAL = 5000
def define_network(arch, batch_size, sample_factor, crop_size=None):
"""Creates the network architecture.
Args:
arch: The architecture type. "frrn_a" or "frrn_b"
batch_size: The batch size.
sample_factor: The subsampling factor.
crop_size: The size of the image crops. None will result in full-frame
training.
Returns:
The newly created network instance.
Raises:
ValueError: If `arch` is not a valid architecture identifier.
"""
if arch == "frrn_a":
builder = architectures.FRRNABuilder
elif arch == "frrn_b":
builder = architectures.FRRNBBuilder
else:
raise ValueError("Invalid network architecture {}.".format(arch))
# Define the theano variables
input_var = T.ftensor4()
builder = builder(
base_channels=BASE_CHANNELS,
lanes=FR_CHANNELS,
multiplier=MULTIPLIER,
num_classes=NUM_CLASSES
)
if crop_size is None:
network = builder.build(
input_var=input_var,
input_shape=(batch_size,
IMAGE_CHANNELS,
IMAGE_ROWS // sample_factor,
IMAGE_COLS // sample_factor))
else:
network = builder.build(
input_var=input_var,
input_shape=(batch_size,
IMAGE_CHANNELS,
crop_size,
crop_size))
return network
def compile_train_function(network, batch_size, learning_rate):
"""Compiles the training function.
Args:
network: The network instance.
batch_size: The training batch size.
learning_rate: The learning rate.
Returns:
The update function that takes a batch of images and targets and updates the
network weights.
"""
learning_rate = np.float32(learning_rate)
input_var = network.input_layers[0].input_var
target_var = T.ftensor4()
# Loss function
loss_fn = functools.partial(
losses.bootstrapped_xentropy,
targets=target_var,
batch_size=batch_size,
multiplier=BOOTSTRAP_MULTIPLIER
)
# Update function
lr = theano.shared(learning_rate)
update_fn = functools.partial(lasagne.updates.adam, learning_rate=lr)
pylogging.info("Compile SGD updates")
gd_step = hybrid_training.compile_gd_step(
network, loss_fn, [input_var, target_var], update_fn)
reduce_lr = theano.function(
inputs=[],
updates=collections.OrderedDict([
(lr, T.maximum(np.float32(5e-5), lr / np.float32(1.25)))
])
)
def _compute_update(imgs, targets, update_counter):
if (update_counter + 1) % REDUCE_LR_INTERVAL == 0:
reduce_lr()
loss = gd_step(imgs, targets)
return loss
return _compute_update
def compile_validation_function(network, batch_size):
"""Compiles the validation function.
Args:
network: The network instance.
batch_size: The batch size.
Returns:
A function that takes in a batch of images and targets and returns the
predicted segmentation mask and the loss.
"""
input_var = network.input_layers[0].input_var
target_var = T.ftensor4()
predictions = lasagne.layers.get_output(
network.output_layers, deterministic=True)[0]
loss = losses.bootstrapped_xentropy(
predictions=predictions,
targets=target_var,
batch_size=batch_size,
multiplier=BOOTSTRAP_MULTIPLIER
)
pylogging.info("Compile validation function")
return theano.function(
inputs=[input_var, target_var],
outputs=[T.argmax(predictions, axis=1), loss]
)
def get_training_provider(cityscapes_folder,
sample_factor,
batch_size,
iterator_type,
crop_size):
"""Creates the training data provider.
Args:
cityscapes_folder: The folder in which the Cityscapes Dataset is
located.
sample_factor: The image sampling factor.
batch_size: The batch size.
iterator_type: The iterator type. "uniform" or "weighted".
crop_size: The size of the image crops. None will result in full-frame
training.
Returns:
A chianti data provider.
"""
augmentors = [
pychianti.Augmentor.Translation(120),
]
if sample_factor > 1:
augmentors.append(pychianti.Augmentor.Subsample(sample_factor))
if crop_size is not None:
augmentors.append(pychianti.Augmentor.Crop(crop_size, NUM_CLASSES))
augmentors.extend([
pychianti.Augmentor.Gamma(0.05),
pychianti.Augmentor.Rotation(10),
pychianti.Augmentor.Saturation(0.5, 1.5),
pychianti.Augmentor.Hue(-30, 30),
])
images = utility.get_image_label_pairs(cityscapes_folder, "train")
if iterator_type == "uniform":
iterator = pychianti.Iterator.Random(images)
elif iterator_type == "weighted":
# Load the image weights
with open("data_weights.pkl", "rb") as f:
w = pickle.load(f)
weights = []
for img in images:
image_name = img[0].split("/")[-1]
weights.append(w[image_name])
iterator = pychianti.Iterator.WeightedRandom(images, weights)
else:
raise ValueError("Invalid iterator type {}.".format(iterator_type))
provider = pychianti.DataProvider(
pychianti.Augmentor.Combined(augmentors),
pychianti.Loader.RGB(),
pychianti.Loader.ValueMapper(utility.cityscapes_value_map),
iterator,
batch_size,
NUM_CLASSES)
return provider
def get_validation_provider(cityscapes_folder,
sample_factor,
batch_size,
crop_size=None):
"""Creates the validation data provider.
Args:
cityscapes_folder: The folder in which the Cityscapes Dataset is
located.
sample_factor: The image sampling factor.
batch_size: The batch size.
crop_size: The size of the image crops. None will result in full-frame
training.
Returns:
A chianti data provider.
"""
augmentors = []
if sample_factor > 1:
augmentors.append(pychianti.Augmentor.Subsample(sample_factor))
if crop_size is not None:
augmentors.append(pychianti.Augmentor.Crop(crop_size, NUM_CLASSES))
validation_images = utility.get_image_label_pairs(cityscapes_folder, "val")
return pychianti.DataProvider(
pychianti.Augmentor.Combined(augmentors),
pychianti.Loader.RGB(),
pychianti.Loader.ValueMapper(utility.cityscapes_value_map),
pychianti.Iterator.Sequential(validation_images),
batch_size,
NUM_CLASSES)
def main():
"""Trains a FRRN architecture on the Cityscapes Dataset."""
parser = argparse.ArgumentParser(
description="Trains a Full-Resolution Residual"
" Network on the Cityscapes"
" Dataset.")
parser.add_argument("--architecture",
type=str,
choices=["frrn_a", "frrn_b"],
required=True,
help="The network architecture type.")
parser.add_argument("--model_file",
type=str,
required=True,
help="The model filename. Weights are initialized to "
"the given values if the file exists. Snapshots "
"are stored using a _snapshot_[iteration] "
"post-fix.")
parser.add_argument("--log_file",
type=str,
required=True,
help="The log filename. Use log_monitor.py in order to "
"monitor training progress in the terminal.")
parser.add_argument("--cs_folder",
type=str,
required=True,
help="The folder that contains the Cityscapes Dataset.")
parser.add_argument("--batch_size",
type=int,
default=3,
help="The batch size.")
parser.add_argument("--validation_interval",
type=int,
default=500,
help="The validation interval.")
parser.add_argument("--iterator",
type=str,
default="uniform",
choices=["uniform", "weighted"],
help="The dataset iterator type.")
parser.add_argument("--crop_size",
type=int,
default=0,
help="The size of crops to extract from the "
"full-resolution images. If 0, then no crops "
"will be extracted.")
parser.add_argument("--learning_rate",
type=float,
default=1e-3,
help="The learning rate to use.")
parser.add_argument("--sample_factor",
type=int,
default=0,
help="The sampling factor.")
args = parser.parse_args()
# Determine the sampling factor based on the network architecture
if args.architecture == "frrn_a":
sample_factor = 4
else:
sample_factor = 2
if args.sample_factor != 0:
sample_factor = args.sample_factor
if args.crop_size > 0:
crop_size = args.crop_size
sample_factor = 1
else:
crop_size = None
pylogging.info("Sample factor: {}".format(sample_factor))
# Define the network lasagne graph and try to load the model file
network = define_network(args.architecture,
args.batch_size,
sample_factor,
crop_size)
pylogging.info("Try to load weights from {}".format(args.model_file))
network.load_model(args.model_file)
# Get the logger
logger = logging.FileLogWriter(args.log_file)
# Create the optimizer
opt = optimizer.MiniBatchOptimizer(
compile_train_function(network, args.batch_size, args.learning_rate),
get_training_provider(args.cs_folder,
sample_factor,
args.batch_size,
args.iterator,
crop_size),
[
hooks.SnapshotHook(
args.model_file,
network,
interval=args.validation_interval),
hooks.LoggingHook(logger),
hooks.SegmentationValidationHook(
compile_validation_function(network, args.batch_size),
get_validation_provider(args.cs_folder,
sample_factor,
args.batch_size,
crop_size),
logger,
interval=args.validation_interval)
])
pylogging.info("Start training")
opt.optimize()
if __name__ == "__main__":
pylogging.basicConfig(format="%(asctime)s %(levelname)s %(message)s",
level=pylogging.DEBUG)
main()
