Class Attribute_Classifier2

model/vision_transformer.py:408–467

class Attribute_Classifier2(nn.Module):
    def __init__(self, dict_attribute, in_dim, projected_dim, out_dim, norm_type='bn'):
        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

        for key in dict_attribute.keys():
            individual_head = nn.ModuleList()

            embedding_layers = [nn.Linear(in_dim, projected_dim)]
            if use_norm:
                embedding_layers.append(norm_class(projected_dim))
            embedding_layers.append(nn.GELU())
            embedding_layers = nn.Sequential(*embedding_layers)
            output_layers = [nn.Linear(projected_dim, out_dim)]
            if use_norm:
                output_layers.append(norm_class(out_dim))
            output_layers = nn.Sequential(*output_layers)
            individual_head.append(embedding_layers)
            individual_head.append(output_layers)
            individual_head.append(nn.Linear(out_dim, len(dict_attribute[key]) + 1)) # 1 for no present

            self.classifier_list.append(individual_head)

        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 = []
        attribute_embedding = []
        for individual_head in self.classifier_list:
            projected = individual_head[0](x)
            _attribute_embedding = individual_head[1](projected)
            attribute_embedding.append(_attribute_embedding.detach().clone())
            logit = individual_head[2](_attribute_embedding)
            probability.append(self._log_softmax(logit))
        attribute_embedding = torch.cat(attribute_embedding, dim=1)

        if self.training == True:
            return probability