Testing ResNet Components

res_net_block.py

import torch
import torch.nn as nn

class ResNetBlock(nn.Module):
    def __init__(self, input_channel_count, internal_channel_count, downsampler = None, stride=1, layers = 'all'):
        # input_channel_count:  An image tensor is a h*w*k tensor where k is represents
        #    how many channels, naively, groups of features discovered by the ml model.
        #    the shape of the input tensor is h*w*input_channel.
        # internal_channel_count:  like above, but the shape of the output tensor is
        #    h*w*internal_channel_count.  In a typical convnet strategy, you will want the
        #    internal_channel_count to grow with each successive layer application.
        # downsampler:  A shape-changing identity layer that is necessary if we change
        #    the input size or if we change the number of channels.
        # stride: unless 1, skips (stride - 1) pixels in the image as a downsizing strategy.

        super(ResNetBlock, self).__init__()
        # useful for testing.
        self.use_layers = layers

        # a parameter which expresses the multiplicity of the channels coming out relative
        # the channels coming in.
        self.expansion = 4
        self.conv1 = nn.Conv2d(input_channel_count, internal_channel_count, kernel_size=1, stride=1, padding=0)
        self.batchnorm1 = nn.BatchNorm2d(internal_channel_count)
        self.conv2 = nn.Conv2d(internal_channel_count, internal_channel_count, kernel_size=3, stride=stride, padding=1)
        self.batchnorm2 = nn.BatchNorm2d(internal_channel_count)
        self.conv3 = nn.Conv2d(internal_channel_count, internal_channel_count * self.expansion, kernel_size=1, stride=1, padding=0)
        self.batchnorm3 = nn.BatchNorm2d(internal_channel_count * self.expansion)
        self.relu = nn.ReLU()
        self.downsampler = downsampler

    def forward(self, activation):
        identity = activation.clone()

        # apply the predefined layers, generating activations as you go along.
        activation = self._apply_layer(activation, self.conv1, 1)
        activation = self._apply_layer(activation, self.batchnorm1, 2)
        activation = self._apply_layer(activation, self.relu, 3)
        activation = self._apply_layer(activation, self.conv2, 4)
        activation = self._apply_layer(activation, self.batchnorm2, 5)
        activation = self._apply_layer(activation, self.relu, 6)
        activation = self._apply_layer(activation, self.conv3, 7)
        activation = self._apply_layer(activation, self.batchnorm3, 8)

        # we may have to alter the shape of the initial input
        if self.downsampler is not None:
            identity = self.downsampler(identity)

        activation = self._apply_layer(activation, lambda a: a + identity, 9)
        output = self._apply_layer(activation, self.relu, 10)
        return output

    """
    provides a downsampler neural net for general use.  This can be called,
    because the calling method is overloaded.

    Note this is a static function.
    """
    def downsampler_fn(input_channel_count, internal_channel_count, stride):
        return nn.Sequential(
                    nn.Conv2d(
                        input_channel_count,
                        internal_channel_count * 4,
                        kernel_size=1,
                        stride=stride,
                    ),
                    nn.BatchNorm2d(internal_channel_count * 4))

    def _apply_layer(self, activation, layer, layer_id):
        if self.use_layers == 'all' or layer_id <= self.use_layers:
            return layer(activation)
        else:
            return activation

test_res_net_block.py

import unittest
import torch
from res_net_block import ResNetBlock

def GenBlock(layers, downsampler = None):
    return ResNetBlock(64, 64, layers = layers, downsampler = downsampler)

# torch tensors are row-major.  Makes you want to pull your teeth out.
# the order is as follows:  batch index, channels, image column, image row.
INPUT_TENSOR = torch.randn(4, 64, 224, 224)
class TestResNetBlock(unittest.TestCase):
    def test_layer_1(self):
        net = GenBlock(1)
        # layer 1 (conv1) doesn't change the tensor size.
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 64, 224, 224])

    def test_layer_2(self):
        net = GenBlock(2)
        # layer 2 (batchnorm1) doesn't change the tensor size.
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 64, 224, 224])

    def test_layer_3(self):
        net = GenBlock(3)
        # layer 3 (relu) doesn't change the tensor size.
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 64, 224, 224])

    def test_layer_4(self):
        net = GenBlock(4)
        # layer 4 (conv2) doesn't change the tensor size.
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 64, 224, 224])

    def test_layer_5(self):
        net = GenBlock(5)
        # layer 5 (batchnorm2) doesn't change the tensor size.
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 64, 224, 224])

    def test_layer_6(self):
        net = GenBlock(6)
        # layer 6 (relu) doesn't change the tensor size.
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 64, 224, 224])

    def test_layer_7(self):
        net = GenBlock(7)
        # layer 7 (conv3) makes there be 4x more channels
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 256, 224, 224])

    def test_layer_8(self):
        # layer 8 (batchnorm3) doesn't change the tensor size.
        net = GenBlock(8)
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 256, 224, 224])

    def test_layer_9(self):
        # layer 9 (addition origin) in the basic block case, must
        # have a supplied downsampler
        downsampler = ResNetBlock.downsampler_fn(64, 64, 1)
        net = GenBlock(9, downsampler = downsampler)
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 256, 224, 224])

    def test_layer_10(self):
        # layer 10 (addition origin) in the basic block case, must
        # have a supplied downsampler
        downsampler = ResNetBlock.downsampler_fn(64, 64, 1)
        net = GenBlock(10, downsampler = downsampler)
        output = net(INPUT_TENSOR).to("cpu")
        self._assertSize(output, [4, 256, 224, 224])

    def _assertSize(self, tensor, sizespec):
        self.assertEqual(tensor.size(), torch.Size(sizespec))