def __init__(self, name, n_classes=NUM_CLASSES):
super().__init__()
base = densenet(name=name, n_classes=n_classes)
feat_dim = base.last_linear.in_features
self.base = base
self.pool = GlobalConcatPool2d()
self.head = nn.Sequential(nn.Linear(feat_dim * 2, feat_dim * 2),
nn.BatchNorm1d(feat_dim * 2),
nn.ReLU(inplace=True), nn.Dropout(0.25),
nn.Linear(feat_dim * 2, n_classes))
def __init__(self, model_name, num_classes):
super().__init__()
self.m1 = create_model(model_name)
self.m2 = create_model(model_name)
self.pool = GlobalConcatPool2d()
out_features = self.m1.last_linear.in_features * 2
self.top = nn.Sequential(
nn.Linear(out_features * 2, out_features),
nn.LeakyReLU(),
nn.Dropout(),
nn.Linear(out_features, num_classes)
)