Deep-Learning-with-PyTorch

Updating the model for segmentation
367
probability for every pixel: that is, performing segmentation. Rather than implementing
a custom U-Net segmentation model from scratch, we’re going to appropriate
an existing implementation from an open source repository on GitHub.
The U-Net implementation at https://github.com/jvanvugt/pytorch-unet seems
to meet our needs well. 4 It’s MIT licensed (copyright 2018 Joris), it’s contained in a
single file, and it has a number of parameter options for us to tweak. The file is
included in our code repository at util/unet.py, along with a link to the original repository
and the full text of the license used.
NOTE While it’s less of an issue for personal projects, it’s important to be
aware of the license terms attached to open source software you use for a project.
The MIT license is one of the most permissive open source licenses, and
it still places requirements on users of MIT licensed code! Also be aware that
authors retain copyright even if they publish their work in a public forum
(yes, even on GitHub), and if they do not include a license, that does not
mean the work is in the public domain. Quite the opposite! It means you
don’t have any license to use the code, any more than you’d have the right to
wholesale copy a book you borrowed from the library.
We suggest taking some time to inspect the code and, based on the knowledge you have
built up until this point, identify the building blocks of the architecture as they are
reflected in the code. Can you spot skip connections? A particularly worthy exercise for
you is to draw a diagram that shows how the model is laid out, just by looking at the code.
Now that we have found a U-Net implementation that fits the bill, we need to
adapt it so that it works well for our needs. In general, it’s a good idea to keep an eye
out for situations where we can use something off the shelf. It’s important to have a
sense of what models exist, how they’re implemented and trained, and whether any
parts can be scavenged and applied to the project we’re working on at any given
moment. While that broader knowledge is something that comes with time and experience,
it’s a good idea to start building that toolbox now.
13.4.1 Adapting an off-the-shelf model to our project
We will now make some changes to the classic U-Net, justifying them along the way. A
useful exercise for you will be to compare results between the vanilla model and the
one after the tweaks, preferably removing one at a time to see the effect of each
change (this is also called an ablation study in research circles).
First, we’re going to pass the input through batch normalization. This way, we
won’t have to normalize the data ourselves in the dataset; and, more importantly, we
will get normalization statistics (read mean and standard deviation) estimated over
individual batches. This means when a batch is dull for some reason—that is, when
there is nothing to see in all the CT crops fed into the network—it will be scaled more
4
The implementation included here differs from the official paper by using average pooling instead of max
pooling to downsample. The most recent version on GitHub has changed to use max pool.