| 7 | |
| 8 | class Attribute_Classifier(nn.Module): |
| 9 | def __init__(self, dict_attribute, in_dim, projected_dim, use_bn=False, use_independent_projection=True): |
| 10 | super().__init__() |
| 11 | print(dict_attribute) |
| 12 | self.num_attribute_class = len(dict_attribute.keys()) |
| 13 | self.num_attribute_all = sum([len(v) for v in dict_attribute.values()]) |
| 14 | self.classifier_list = nn.ModuleList() |
| 15 | self._softmax = nn.Softmax() |
| 16 | if use_independent_projection: |
| 17 | self.shared_projected_layer = None |
| 18 | else: |
| 19 | if use_bn: |
| 20 | self.shared_projected_layer = nn.Sequential(nn.Linear(in_dim, projected_dim), |
| 21 | nn.BatchNorm1d(projected_dim), |
| 22 | nn.GELU()) |
| 23 | else: |
| 24 | self.shared_projected_layer = nn.Sequential(nn.Linear(in_dim, projected_dim), |
| 25 | nn.GELU()) |
| 26 | for key in dict_attribute.keys(): |
| 27 | if use_independent_projection: |
| 28 | layers = [nn.Linear(in_dim, projected_dim)] |
| 29 | if use_bn: |
| 30 | layers.append(nn.BatchNorm1d(projected_dim)) |
| 31 | layers.append(nn.GELU()) |
| 32 | layers.append(nn.Linear(projected_dim, len(dict_attribute[key]) + 1)) # 1 for no present |
| 33 | self.classifier_list.append(nn.Sequential(*layers)) |
| 34 | else: |
| 35 | layers = [] |
| 36 | if use_bn: |
| 37 | layers.append(nn.BatchNorm1d(projected_dim)) |
| 38 | layers.append(nn.GELU()) |
| 39 | layers.append(nn.Linear(projected_dim, len(dict_attribute[key]) + 1)) # 1 for no present |
| 40 | self.classifier_list.append(nn.Sequential(*layers)) |
| 41 | self.apply(self._init_weights) |
| 42 | |
| 43 | def _init_weights(self, m): |
| 44 | if isinstance(m, nn.Linear): |