Class Meta_Attribute_Generator2

model/attribute_transformer.py:881–911

class Meta_Attribute_Generator2(nn.Module):
    def __init__(self, vit_backbone_model, dict_attribute, grad_from_block=11):
        super().__init__()
        backbone_feat_dim = vit_backbone_model.num_features
        self.num_attribute_class = len(dict_attribute.keys())
        self.num_attribute_all = sum([len(v) for v in dict_attribute.values()])
        self.attribute_generator_list = nn.ModuleList()

        for key in dict_attribute.keys():
            _conv = attribute_subnet(backbone_feat_dim)
            _classifier = nn.Linear(backbone_feat_dim, len(dict_attribute[key]) + 1)  # 1 for no present
            _softmax = nn.Softmax(dim=1)
            self.attribute_generator_list.append(nn.Sequential(_conv, _classifier, _softmax))

        del vit_backbone_model
        torch.cuda.empty_cache()
        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):
        fake_prob_list = []
        meta_embedding = torch.transpose(x, 1, 2)
        for att_head in self.attribute_generator_list:
            fake_prob = att_head(meta_embedding)
            fake_prob_list.append(fake_prob)
        return fake_prob_list


def at_small(pretrain_path):