Class Attribute_Classifier

model/vision_transformer.py:330–404

class Attribute_Classifier(nn.Module):
    def __init__(self, dict_attribute, in_dim, projected_dim, norm_type='bn', 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._log_softmax = nn.LogSoftmax(dim=1)
        if norm_type is None or norm_type == 'none':
            use_norm = False
        elif norm_type == 'bn':
            norm_class = nn.BatchNorm1d
            use_norm = True
        elif norm_type == 'ln':
            norm_class = nn.LayerNorm
            use_norm = True
        else:
            raise NotImplementedError

        if use_independent_projection:
            self.shared_projected_layer = None
        else:
            if use_norm:
                self.shared_projected_layer = nn.Sequential(nn.Linear(in_dim, projected_dim * 2),
                                                            norm_class(projected_dim * 2),
                                                            nn.GELU(),
                                                            nn.Linear(projected_dim * 2, projected_dim),
                                                            norm_class(projected_dim),
                                                            )
            else:
                self.shared_projected_layer = nn.Sequential(nn.Linear(in_dim, projected_dim * 2),
                                                            nn.GELU(),
                                                            nn.Linear(projected_dim * 2, projected_dim)
                                                            )

        for key in dict_attribute.keys():
            if use_independent_projection:
                layers = [nn.Linear(in_dim, projected_dim)]
                if use_norm:
                    layers.append(norm_class(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:
                _classifier = nn.Linear(projected_dim, len(dict_attribute[key]) + 1)  # 1 for no present
                self.classifier_list.append(_classifier)
        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:
            attribute_embedding = []