models.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch


class Flatten(torch.nn.Module):

    def forward(self, x):

        return x.view(x.size(0), -1)


def get_cnn4(num_classes):

    def convbn(inplanes, planes):
        return torch.nn.Sequential(
            torch.nn.Conv2d(inplanes, planes, 3, padding=1),
            torch.nn.BatchNorm2d(planes),
            torch.nn.ReLU(True),
            torch.nn.MaxPool2d(2)
        )

    class CNN4(torch.nn.Module):

        def __init__(self, num_classes):

            super(CNN4, self).__init__()
            self.features = torch.nn.Sequential(
                convbn(3, 64),
                convbn(64, 64),
                convbn(64, 64),
                convbn(64, 64),
                Flatten()
            )
            self.head = torch.nn.Linear(1600, num_classes)

        def forward(self, x):

            x = self.features(x)
            x = self.head(x)
            return x

    return CNN4(num_classes)


def get_resnet12(num_classes):

    class BasicBlock(torch.nn.Module):

        def __init__(self, inplanes, planes):

            super(BasicBlock, self).__init__()

            def conv3x3(in_planes, out_planes):
                return torch.nn.Conv2d(in_planes, out_planes, kernel_size=3,
                                       padding=1, bias=False)

            self.conv1 = conv3x3(inplanes, planes)
            self.bn1 = torch.nn.BatchNorm2d(planes)
            self.conv2 = conv3x3(planes, planes)
            self.bn2 = torch.nn.BatchNorm2d(planes)
            self.conv3 = conv3x3(planes, planes)
            self.bn3 = torch.nn.BatchNorm2d(planes)
            self.maxpool = torch.nn.MaxPool2d(2)
            self.downsample = torch.nn.Sequential(
                torch.nn.Conv2d(inplanes, planes, kernel_size=1, bias=False),
                torch.nn.BatchNorm2d(planes),
            )
            self.relu = torch.nn.ReLU(True)

        def forward(self, x):

            residual = self.downsample(x)

            out = self.conv1(x)
            out = self.bn1(out)
            out = self.relu(out)

            out = self.conv2(out)
            out = self.bn2(out)
            out = self.relu(out)

            out = self.conv3(out)
            out = self.bn3(out)

            out += residual
            out = self.relu(out)
            out = self.maxpool(out)

            return out

    class ResNet12(torch.nn.Module):

        def __init__(self, num_classese):

            super(ResNet12, self).__init__()
            self.features = torch.nn.Sequential(
                BasicBlock(3, 64),
                BasicBlock(64, 128),
                BasicBlock(128, 256),
                BasicBlock(256, 512),
                torch.nn.AvgPool2d(5),
                Flatten()
            )
            self.head = torch.nn.Linear(512, num_classes)

            for m in self.modules():
                if isinstance(m, torch.nn.Conv2d):
                    torch.nn.init.kaiming_normal_(m.weight, mode='fan_out',
                                                  nonlinearity='relu')

        def forward(self, x):

            x = self.features(x)
            x = self.head(x)
            return x

    return ResNet12(num_classes)


MODELS = {
    'cnn4': get_cnn4,
    'resnet12': get_resnet12
}