def build_model()

in models/__init__.py [0:0]


def build_model(params):
    if params.architecture == "lenet":
        return LeNet(params)
    elif params.architecture == "kllenet":
        return KLLeNet(params)
    elif params.architecture == "linear":
        return LinearNet(params)
    elif params.architecture == "mlp":
        return MLP(params)
    elif params.architecture=="alexnet":
        return AlexNet(params)
    elif params.architecture == "lstm":
        return LSTMLM(params)
    elif params.architecture == "resnet18":
        return models.resnet18(pretrained=False)
    elif params.architecture == "smallnet":
        return nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(start_dim=1, end_dim=-1),
            nn.Linear(128, params.num_classes, bias=True),
        )
    elif params.architecture == "leaks":
        return nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(start_dim=1, end_dim=-1),
            nn.Linear(128, params.num_classes, bias=True),
        )