model.py

def get_senet50(config):
    model = se_resnext50_32x4d(pretrained='imagenet')
#     print(model.last_linear)
#     print(model.dropout)
    model.layer0.conv1 = nn.Conv2d(config.channels, 64, 3, stride=2, padding=1, bias=False)
    model.avg_pool = nn.AdaptiveAvgPool2d(1)
    model.dropout = None
    model.last_linear = nn.Sequential(
        nn.BatchNorm1d(2048),
        nn.Dropout(0.5),
        nn.Linear(2048, config.num_classes),
    )
    return model

def get_senet101(config):
    model = se_resnext101_32x4d(pretrained='imagenet')
#     print(model.last_linear)
#     print(model.dropout)
    model.layer0.conv1 = nn.Conv2d(config.channels, 64, 3, stride=2, padding=1, bias=False)
    model.avg_pool = nn.AdaptiveAvgPool2d(1)
    model.dropout = None
    model.last_linear = nn.Sequential(
        nn.BatchNorm1d(2048),
        nn.Dropout(0.5),
        nn.Linear(2048, config.num_classes),
    )
    return model