| 455 | |
| 456 | |
| 457 | class attribute_subnet(nn.Module): |
| 458 | def __init__(self, input_feature_dim, norm_type='bn'): |
| 459 | super().__init__() |
| 460 | self.conv_1_1 = nn.Conv1d(input_feature_dim, input_feature_dim, 1) |
| 461 | if norm_type == 'bn': |
| 462 | self.norm1 = nn.BatchNorm1d(input_feature_dim) |
| 463 | elif norm_type == 'ln': |
| 464 | self.norm1 = nn.LayerNorm(input_feature_dim) |
| 465 | elif norm_type == 'none' or norm_type is None: |
| 466 | self.norm1 == nn.Identity() |
| 467 | else: |
| 468 | raise NotImplementedError |
| 469 | self.activation = nn.GELU() |
| 470 | self.pool = nn.AdaptiveMaxPool1d(1) |
| 471 | self.flatten = nn.Flatten() |
| 472 | self.norm2 = nn.LayerNorm(input_feature_dim) |
| 473 | |
| 474 | def forward(self, x): |
| 475 | out = self.flatten(self.pool(self.activation(self.norm1(self.conv_1_1(x))))) |
| 476 | out = self.norm2(out) |
| 477 | return out |
| 478 | |
| 479 | |
| 480 | class AttributeTransformer8(nn.Module): |