| 407 | |
| 408 | class Attribute_Classifier2(nn.Module): |
| 409 | def __init__(self, dict_attribute, in_dim, projected_dim, out_dim, norm_type='bn'): |
| 410 | super().__init__() |
| 411 | print(dict_attribute) |
| 412 | self.num_attribute_class = len(dict_attribute.keys()) |
| 413 | self.num_attribute_all = sum([len(v) for v in dict_attribute.values()]) |
| 414 | self.classifier_list = nn.ModuleList() |
| 415 | self._log_softmax = nn.LogSoftmax(dim=1) |
| 416 | if norm_type is None or norm_type == 'none': |
| 417 | use_norm = False |
| 418 | elif norm_type == 'bn': |
| 419 | norm_class = nn.BatchNorm1d |
| 420 | use_norm = True |
| 421 | elif norm_type == 'ln': |
| 422 | norm_class = nn.LayerNorm |
| 423 | use_norm = True |
| 424 | else: |
| 425 | raise NotImplementedError |
| 426 | |
| 427 | for key in dict_attribute.keys(): |
| 428 | individual_head = nn.ModuleList() |
| 429 | |
| 430 | embedding_layers = [nn.Linear(in_dim, projected_dim)] |
| 431 | if use_norm: |
| 432 | embedding_layers.append(norm_class(projected_dim)) |
| 433 | embedding_layers.append(nn.GELU()) |
| 434 | embedding_layers = nn.Sequential(*embedding_layers) |
| 435 | output_layers = [nn.Linear(projected_dim, out_dim)] |
| 436 | if use_norm: |
| 437 | output_layers.append(norm_class(out_dim)) |
| 438 | output_layers = nn.Sequential(*output_layers) |
| 439 | individual_head.append(embedding_layers) |
| 440 | individual_head.append(output_layers) |
| 441 | individual_head.append(nn.Linear(out_dim, len(dict_attribute[key]) + 1)) # 1 for no present |
| 442 | |
| 443 | self.classifier_list.append(individual_head) |
| 444 | |
| 445 | self.apply(self._init_weights) |
| 446 | |
| 447 | def _init_weights(self, m): |
| 448 | if isinstance(m, nn.Linear): |