Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| biases = False | |
| class Pool2BN(nn.Module): | |
| def __init__(self, num_channels): | |
| super().__init__() | |
| self.bn = torch.nn.BatchNorm1d(num_channels * 2) | |
| def forward(self, x): | |
| avgp = torch.nn.functional.adaptive_avg_pool1d(x, 1)[:, :, 0] | |
| maxp = torch.nn.functional.adaptive_max_pool1d(x, 1)[:, :, 0] | |
| x = torch.cat((avgp, maxp), axis=1) | |
| x = self.bn(x) | |
| return x | |
| class MLP(torch.nn.Module): | |
| def __init__(self, layer_sizes, biases=False, sigmoid=False, dropout=None): | |
| super().__init__() | |
| layers = [] | |
| prev_size = layer_sizes[0] | |
| for i, s in enumerate(layer_sizes[1:]): | |
| if i != 0 and dropout is not None: | |
| layers.append(torch.nn.Dropout(dropout)) | |
| layers.append(torch.nn.Linear(in_features=prev_size, out_features=s, bias=biases)) | |
| if i != len(layer_sizes) - 2: | |
| if sigmoid: | |
| # layers.append(torch.nn.Sigmoid()) | |
| layers.append(torch.nn.Tanh()) | |
| else: | |
| layers.append(torch.nn.ReLU()) | |
| layers.append(torch.nn.BatchNorm1d(s)) | |
| prev_size = s | |
| self.mlp = torch.nn.Sequential(*layers) | |
| def forward(self, x): | |
| return self.mlp(x) | |
| class SimpleCNN(torch.nn.Module): | |
| def __init__(self, k, num_filters, sigmoid=False, additional_layer=False): | |
| super(SimpleCNN, self).__init__() | |
| self.sigmoid = sigmoid | |
| self.cnn = torch.nn.Conv1d(in_channels=4, out_channels=num_filters, kernel_size=k, bias=biases) | |
| self.additional_layer = additional_layer | |
| if additional_layer: | |
| self.bn = nn.BatchNorm1d(num_filters) | |
| # self.do = nn.Dropout(0.5) | |
| self.cnn2 = nn.Conv1d(in_channels=num_filters, out_channels=num_filters, kernel_size=1, bias=biases) | |
| self.post = Pool2BN(num_filters) | |
| def forward(self, x): | |
| x = self.cnn(x) | |
| x = (torch.tanh if self.sigmoid else torch.relu)(x) | |
| if self.additional_layer: | |
| x = self.bn(x) | |
| # x = self.do(x) | |
| x = self.cnn2(x) | |
| x = (torch.tanh if self.sigmoid else torch.relu)(x) | |
| x = self.post(x) | |
| #print(f'x shape at CNN output: {x.shape}') | |
| return x | |
| class ResNet1dBlock(torch.nn.Module): | |
| def __init__(self, num_filters, k1, internal_filters, k2, dropout=None, dilation=None): | |
| super().__init__() | |
| self.init_do = torch.nn.Dropout(dropout) if dropout is not None else None | |
| self.bn1 = torch.nn.BatchNorm1d(num_filters) | |
| if dilation is None: | |
| dilation = 1 | |
| self.cnn1 = torch.nn.Conv1d(in_channels=num_filters, out_channels=internal_filters, kernel_size=k1, bias=biases, | |
| dilation=dilation, | |
| padding=(k1 // 2) * dilation) | |
| self.bn2 = torch.nn.BatchNorm1d(internal_filters) | |
| self.cnn2 = torch.nn.Conv1d(in_channels=internal_filters, out_channels=num_filters, kernel_size=k2, bias=biases, | |
| padding=k2 // 2) | |
| def forward(self, x): | |
| x_orig = x | |
| x = self.bn1(x) | |
| x = torch.relu(x) | |
| if self.init_do is not None: | |
| x = self.init_do(x) | |
| x = self.cnn1(x) | |
| x = self.bn2(x) | |
| x = torch.relu(x) | |
| x = self.cnn2(x) | |
| return x + x_orig | |
| class ResNet1d(torch.nn.Module): | |
| def __init__(self, num_filters, block_spec, dropout=None, dilation=None): | |
| super().__init__() | |
| blocks = [ResNet1dBlock(num_filters, *spec, dropout=dropout, dilation=dilation) for spec in block_spec] | |
| self.blocks = torch.nn.Sequential(*blocks) | |
| def forward(self, x): | |
| return self.blocks(x) | |
| class LogisticRegressionTorch(nn.Module): | |
| def __init__(self, input_dim: int, output_dim: int): | |
| super(LogisticRegressionTorch, self).__init__() | |
| self.batch_norm = nn.BatchNorm1d(num_features=input_dim) | |
| self.linear = nn.Linear(input_dim, output_dim) | |
| def forward(self, x): | |
| x = self.batch_norm(x) | |
| out = self.linear(x) | |
| return out | |
| class BertClassifier(nn.Module): | |
| def __init__(self, bert_model: AutoModel, classifier: LogisticRegressionTorch, num_labels: int): | |
| super(BertClassifier, self).__init__() | |
| self.bert = bert_model | |
| self.classifier = classifier | |
| self.num_labels = num_labels | |
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor = None): | |
| outputs = self.bert(input_ids, attention_mask=attention_mask, output_hidden_states=True) | |
| pooled_output = outputs.hidden_states[-1][:, 0, :] | |
| logits = self.classifier(pooled_output) | |
| return logits |