Class Attribute_Classifier

model/attribute_classifier.py:8–61

from torch.nn.init import trunc_normal_

class Attribute_Classifier(nn.Module):
    def __init__(self, dict_attribute, in_dim, projected_dim, use_bn=False, use_independent_projection=True):
        super().__init__()
        print(dict_attribute)
        self.num_attribute_class = len(dict_attribute.keys())
        self.num_attribute_all = sum([len(v) for v in dict_attribute.values()])
        self.classifier_list = nn.ModuleList()
        self._softmax = nn.Softmax()
        if use_independent_projection:
            self.shared_projected_layer = None
        else:
            if use_bn:
                self.shared_projected_layer = nn.Sequential(nn.Linear(in_dim, projected_dim),
                                                            nn.BatchNorm1d(projected_dim),
                                                            nn.GELU())
            else:
                self.shared_projected_layer = nn.Sequential(nn.Linear(in_dim, projected_dim),
                                                            nn.GELU())
        for key in dict_attribute.keys():
            if use_independent_projection:
                layers = [nn.Linear(in_dim, projected_dim)]
                if use_bn:
                    layers.append(nn.BatchNorm1d(projected_dim))
                layers.append(nn.GELU())
                layers.append(nn.Linear(projected_dim, len(dict_attribute[key]) + 1)) # 1 for no present
                self.classifier_list.append(nn.Sequential(*layers))
            else:
                layers = []
                if use_bn:
                    layers.append(nn.BatchNorm1d(projected_dim))
                layers.append(nn.GELU())
                layers.append(nn.Linear(projected_dim, len(dict_attribute[key]) + 1))  # 1 for no present
                self.classifier_list.append(nn.Sequential(*layers))
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        probability = []
        if self.shared_projected_layer is None:
            for classifier in self.classifier_list:
                logit = classifier(x)
                probability.append(self._softmax(logit))
        else:
            x = self.shared_projected_layer(x)
            for classifier in self.classifier_list:
                logit = classifier(x)
                probability.append(self._softmax(logit))

        return probability


class Mlp(nn.Module):