【暂存】multi-head output

class ConResNet(nn.Module):
    """backbone + projection head"""
    def __init__(self, name='resnet50', head='mlp', feat_dim=128, selfcon_pos=[False,False,False], selfcon_arch='resnet', selfcon_size='same', dataset=''):
        super(ConResNet, self).__init__()
        model_fun, dim_in = model_dict[name]
        self.encoder = model_fun(selfcon_pos=selfcon_pos, selfcon_arch=selfcon_arch, selfcon_size=selfcon_size, dataset=dataset)
        if head == 'linear':
            self.head = nn.Linear(dim_in, feat_dim)
            
            self.sub_heads = []
            for pos in selfcon_pos:
                if pos:
                    self.sub_heads.append(nn.Linear(dim_in, feat_dim))
        elif head == 'mlp':
            self.head = nn.Sequential(
                nn.Linear(dim_in, dim_in),
                nn.ReLU(inplace=True),
                nn.Linear(dim_in, feat_dim)
            )
            
            heads = []
            for pos in selfcon_pos:
                if pos:
                    heads.append(nn.Sequential(
                        nn.Linear(dim_in, dim_in),
                        nn.ReLU(inplace=True),
                        nn.Linear(dim_in, feat_dim)
                    ))
            self.sub_heads = nn.ModuleList(heads)
        else:
            raise NotImplementedError(
                'head not supported: {}'.format(head))

    def forward(self, x):
        sub_feat, feat = self.encoder(x)
        print(sub_feat,feat)
        print("******")
        
        sh_feat = []
        for sf, sub_head in zip(sub_feat, self.sub_heads):
            print(sf,sub_head,sub_feat,self.sub_heads)
            sh_feat.append(F.normalize(sub_head(sf), dim=1))
        print("***")
        print(sh_feat)
        print("*************************")
        input()
        
        feat = F.normalize(self.head(feat), dim=1)
        return sh_feat, feat

Core Code:

sub_feat, feat = self.encoder(x)
print(sub_feat,feat)
print("******")

sh_feat = []
for sf, sub_head in zip(sub_feat, self.sub_heads):
    print(sf,sub_head,sub_feat,self.sub_heads)
    sh_feat.append(F.normalize(sub_head(sf), dim=1))
print("***")
print(sh_feat)
print("*************************")
input()

feat = F.normalize(self.head(feat), dim=1)
return sh_feat, feat

Files already downloaded and verified
[tensor([[0.3709, 0.4060, 0.4357,? ..., 0.4369, 0.7116, 0.6154],
??????? [0.3255, 0.8886, 0.4587,? ..., 0.3097, 0.6247, 0.6515],
??????? [0.3265, 0.5572, 0.5465,? ..., 0.1894, 0.5648, 0.7019],
??????? ...,
??????? [0.3285, 0.7454, 0.3656,? ..., 0.4739, 0.4305, 0.3940],
??????? [4.2099, 2.0158, 2.8063,? ..., 3.5441, 2.7423, 1.8907],
??????? [5.0028, 4.0435, 6.1597,? ..., 6.2176, 2.4172, 1.5914]],
?????? device='cuda:0', grad_fn=<ReshapeAliasBackward0>)] tensor([[0.5324, 0.5882, 0.8503,? ..., 0.7225, 0.8470, 0.3471],
??????? [0.7275, 0.6198, 0.7865,? ..., 0.7751, 0.9618, 0.8199],
??????? [0.7426, 0.7385, 0.9845,? ..., 0.4950, 0.7898, 0.4606],
??????? ...,
??????? [0.6489, 0.7440, 1.0151,? ..., 0.5916, 1.0089, 0.2878],
??????? [0.6675, 1.0148, 1.7714,? ..., 1.4871, 0.8820, 5.0314],
??????? [0.7732, 0.9582, 2.3896,? ..., 1.7296, 0.9024, 5.8410]],
?????? device='cuda:0', grad_fn=<ReshapeAliasBackward0>)
******
tensor([[0.3709, 0.4060, 0.4357,? ..., 0.4369, 0.7116, 0.6154],
??????? [0.3255, 0.8886, 0.4587,? ..., 0.3097, 0.6247, 0.6515],
??????? [0.3265, 0.5572, 0.5465,? ..., 0.1894, 0.5648, 0.7019],
??????? ...,
??????? [0.3285, 0.7454, 0.3656,? ..., 0.4739, 0.4305, 0.3940],
??????? [4.2099, 2.0158, 2.8063,? ..., 3.5441, 2.7423, 1.8907],
??????? [5.0028, 4.0435, 6.1597,? ..., 6.2176, 2.4172, 1.5914]],
?????? device='cuda:0', grad_fn=<ReshapeAliasBackward0>) Sequential(
? (0): Linear(in_features=2048, out_features=2048, bias=True)
? (1): ReLU(inplace=True)
? (2): Linear(in_features=2048, out_features=128, bias=True)
) [tensor([[0.3709, 0.4060, 0.4357,? ..., 0.4369, 0.7116, 0.6154],
??????? [0.3255, 0.8886, 0.4587,? ..., 0.3097, 0.6247, 0.6515],
??????? [0.3265, 0.5572, 0.5465,? ..., 0.1894, 0.5648, 0.7019],
??????? ...,
??????? [0.3285, 0.7454, 0.3656,? ..., 0.4739, 0.4305, 0.3940],
??????? [4.2099, 2.0158, 2.8063,? ..., 3.5441, 2.7423, 1.8907],
??????? [5.0028, 4.0435, 6.1597,? ..., 6.2176, 2.4172, 1.5914]],
?????? device='cuda:0', grad_fn=<ReshapeAliasBackward0>)] ModuleList(
? (0): Sequential(
??? (0): Linear(in_features=2048, out_features=2048, bias=True)
??? (1): ReLU(inplace=True)
??? (2): Linear(in_features=2048, out_features=128, bias=True)
? )
)
***
[tensor([[ 0.0400,? 0.0202,? 0.2806,? ..., -0.0758,? 0.0742, -0.0301],
??????? [ 0.0687,? 0.0602,? 0.2580,? ..., -0.0850,? 0.0880, -0.0687],
??????? [ 0.0607, -0.0127,? 0.2398,? ..., -0.0618,? 0.0845, -0.0610],
??????? ...,
??????? [ 0.0440,? 0.0067,? 0.2463,? ..., -0.0795,? 0.0996, -0.0472],
??????? [ 0.1166,? 0.0371,? 0.1152,? ..., -0.1231, -0.0744,? 0.0168],
??????? [ 0.0630,? 0.0231,? 0.1265,? ..., -0.1107, -0.0084, -0.0108]],
?????? device='cuda:0', grad_fn=<DivBackward0>)]
*************************