| import numpy as np |
| import torch |
| from torch import nn |
| import warnings |
| warnings.filterwarnings("ignore") |
|
|
|
|
| class Transpose(nn.Module): |
| def __init__(self): |
| super(Transpose, self).__init__() |
|
|
| def forward(self, x): |
| return torch.transpose(x, 2, 1) |
|
|
|
|
| """ |
| ------------------------------------------- |
| RnaModel Class Construction |
| ------------------------------------------- |
| """ |
|
|
| class Linear2D(nn.Module): |
| def __init__(self, channel_dim, input_dim, output_dim, inter_dim, only_A=False, flatten_dim=None, input_length=None, add_bias=False, reduction=None, dropout_rate=None): |
| super(Linear2D, self).__init__() |
| self.inter_dim = inter_dim |
| self.flatten_dim = flatten_dim |
| self.only_A = only_A |
| self.reduction = reduction |
| self.dropout_rate = dropout_rate |
| if flatten_dim is not None: |
| self.flatten_proj = nn.Parameter(torch.randn(channel_dim, input_dim, flatten_dim)) |
| self.transA = nn.Parameter(torch.randn(channel_dim, flatten_dim * input_length, inter_dim)) |
| self.transB = nn.Parameter(torch.randn(channel_dim, inter_dim, output_dim)) |
| for param in [self.transA, self.transB, self.flatten_proj]: |
| nn.init.xavier_uniform_(param) |
| elif only_A: |
| self.transA = nn.Parameter(torch.randn(channel_dim, input_dim, output_dim)) |
| nn.init.xavier_uniform_(self.transA) |
| else: |
| self.transA = nn.Parameter(torch.randn(channel_dim, input_dim, inter_dim)) |
| self.transB = nn.Parameter(torch.randn(channel_dim, inter_dim, output_dim)) |
| for param in [self.transA, self.transB]: |
| nn.init.xavier_uniform_(param) |
| self.add_bias = add_bias |
| if add_bias: |
| self.bias = nn.Parameter(torch.zeros(channel_dim, output_dim)) |
| if dropout_rate is not None: |
| self.factor_drop = nn.Dropout1d(dropout_rate) |
|
|
| if reduction is None: |
| self.reduction_fun = lambda x: x |
| elif reduction == 'mean': |
| self.reduction_fun = lambda x: torch.mean(x, 0) |
| elif reduction == 'max': |
| self.reduction_fun = lambda x: torch.max(x, 0)[0] |
| elif reduction == 'mlp': |
| self.channel_factor = nn.Parameter(torch.randn(output_dim, channel_dim, 1)) |
| nn.init.xavier_uniform_(self.channel_factor) |
|
|
| def forward(self, A, is_dropout=True): |
| dimA = A.dim() |
| if is_dropout and self.dropout_rate is not None: |
| transA = self.factor_drop(self.transA) |
| if self.flatten_dim is not None: |
| flatten_proj = self.factor_drop(self.flatten_proj) |
| else: |
| transA = self.transA |
| if self.flatten_dim is not None: |
| flatten_proj = self.flatten_proj |
| if self.flatten_dim is not None: |
| output = torch.einsum('bld,cdr->cblr', A, flatten_proj) |
| output = torch.flatten(output, start_dim=2) |
| output = torch.einsum('cbd,cdr->cbr', output, transA) |
| output = torch.einsum('cbr,cro->cbo', output, self.transB) |
| elif self.only_A: |
| output = torch.einsum('bd,cdr->cbr', A, transA) |
| elif dimA == 3: |
| output = torch.einsum('bld,cdr->cblr', A, transA) |
| output = torch.einsum('cblr,cro->cblo', output, self.transB) |
| elif dimA == 2: |
| output = torch.einsum('bd,cdr->cbr', A, transA) |
| output = torch.einsum('cbr,cro->cbo', output, self.transB) |
| if self.add_bias: |
| bias = self.bias |
| if dimA == 2 or self.flatten_dim is not None or self.only_A: |
| bias = bias.unsqueeze(1) |
| else: |
| bias = bias.unsqueeze(1).unsqueeze(1) |
| output = output + bias |
| if self.reduction == 'mlp': |
| output = torch.einsum('cbd,dcr->bdr', output, self.channel_factor) |
| return output.squeeze(-1) |
| else: |
| return self.reduction_fun(output) |
|
|
| class CNN_Encoder_Model(nn.Module): |
| r""" |
| input shape must be: [batch_size, channel_size, sequence_length] |
| """ |
| def __init__(self, model_config, input_shapes): |
| super(CNN_Encoder_Model, self).__init__() |
| self.model_config = model_config |
| self.input_shapes = input_shapes |
| print("input shapes is {}".format(self.input_shapes)) |
| hyper_use = model_config["training"]["hyper_use"] |
| p = model_config[hyper_use] |
| self.hyperParam = p |
| self.motif_depth = 4 |
| self.bias = nn.Parameter(torch.zeros([1])) |
| self.moving_bias = 0 |
|
|
|
|
| self.detectorDict = nn.ModuleDict() |
| for k in ["reg", "cds"]: |
| self.detectorDict[k] = nn.ModuleList() |
| for _ in range(self.motif_depth): |
| self.detectorDict[k].append( |
| nn.Sequential( |
| nn.Conv1d( |
| in_channels=input_shapes[k][0], |
| out_channels=int(p["filters1"]), |
| kernel_size=int(p["kernel_size1"]), |
| stride=int(p["filters_stride1"]), |
| dilation=int(p["dilation1"]), |
| ), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["filters1"])), |
| ) |
| ) |
|
|
| self.filterDict = nn.ModuleDict() |
| for k in ["reg", "cds"]: |
| self.filterDict[k] = nn.Sequential( |
| nn.BatchNorm1d(int(p["filters1"])), |
| nn.Dropout1d(float(p["dropout1"])), |
| nn.MaxPool1d( |
| kernel_size=int(p["pool_size1"]), |
| stride=int(p["stride1"]), |
| padding=0, |
| ), |
| ) |
|
|
| self.encoderDict = nn.ModuleDict() |
| for k in ["reg", "cds"]: |
| self.encoderDict[k] = nn.Sequential( |
| nn.Conv1d( |
| in_channels=int(p["filters1"]), |
| out_channels=int(p["filters3"]), |
| kernel_size=int(p["kernel_size3"]), |
| stride=int(p["filters_stride3"]), |
| dilation=int(p["dilation3"]), |
| ), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["filters3"])), |
| nn.MaxPool1d( |
| kernel_size=int(p["pool_size3"]), |
| stride=int(p["stride3"]), |
| padding=0, |
| ), |
| Transpose(), |
| nn.Dropout1d(float(p["dropout3"])), |
| |
| ) |
| utr_flatten_size = self.conv_shape("reg") |
|
|
| self.utr_mlp = nn.Sequential( |
| |
| Linear2D(channel_dim=8, input_dim=utr_flatten_size[2], |
| input_length=utr_flatten_size[1], output_dim=int(p['dense4']), |
| inter_dim=4, dropout_rate=0.3, reduction='mlp', |
| flatten_dim=16), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["dense4"])), |
| nn.Dropout(float(p["dropout4"])), |
| nn.Linear(in_features=int(p["dense4"]), out_features=int(p["dense6"])), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["dense6"])), |
| ) |
|
|
| cds_flatten_size = self.conv_shape("cds") |
|
|
| self.RPF_fc = nn.Sequential( |
| |
| Linear2D(channel_dim=8, input_dim=cds_flatten_size[2], |
| input_length=cds_flatten_size[1], output_dim=int(p['dense5']), |
| inter_dim=4, dropout_rate=0.3, reduction='mlp', |
| flatten_dim=16), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["dense5"])), |
| nn.Dropout(float(p["dropout5"])), |
| nn.Linear(in_features=int(p["dense5"]), out_features=1), |
| ) |
|
|
| self.attention_utr = nn.Sequential( |
|
|
| nn.Linear( |
| in_features=int(p["dense6"]), |
| out_features=int(p["filters1"]) * self.motif_depth, |
| ), |
| nn.Tanh(), |
| ) |
|
|
| self.attention_cds = nn.Sequential( |
| nn.Linear( |
| in_features=int(p["dense6"]) + 1, |
| out_features=int(p["filters1"]) * self.motif_depth, |
| ), |
| nn.Tanh(), |
| ) |
|
|
| self.mRNA_layer = nn.Sequential( |
| nn.Linear( |
| in_features=self.input_shapes["trans"][0], out_features=int(p["dense4"]) |
| ), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["dense4"])), |
| nn.Dropout(float(p["dropout4"])), |
| nn.Linear(in_features=int(p["dense4"]), out_features=int(p["dense6"])), |
| nn.ReLU(), |
| nn.BatchNorm1d(int(p["dense6"])), |
| ) |
|
|
| def motif_detection(self, sequence_input, a, seq_type: str): |
|
|
| results = [ |
| motif_detector(sequence_input) |
| for motif_detector in self.detectorDict[seq_type] |
| ] |
| avg_filter = torch.exp(a) |
| features = torch.sum(torch.stack(results, 3) * avg_filter, 3) / torch.sum( |
| avg_filter, 3 |
| ) |
| motif_result = self.filterDict[seq_type](features) |
| return motif_result |
|
|
| def conv_shape(self, seqType: str): |
| x_input_1 = torch.zeros(1, *self.input_shapes[seqType]) |
| x_output = self.motif_detection( |
| x_input_1, torch.zeros([1, 256, 1, self.motif_depth]), seqType |
| ) |
| x_output = self.encoderDict[seqType](x_output) |
| flatten_shape = x_output.size() |
| return flatten_shape |
|
|
| def forward(self, regSequence, cdsSequence, transArray, rnaCount): |
| mRNA_features = self.mRNA_layer(transArray) |
| attention_Reg = torch.reshape( |
| self.attention_utr(mRNA_features), |
| [-1, int(self.hyperParam["filters1"]), 1, self.motif_depth], |
| ) |
| regOutput = self.motif_detection(regSequence, attention_Reg, "reg") |
| regOutput = self.encoderDict["reg"](regOutput) |
| regFeatures = self.utr_mlp(regOutput) |
| attention_cds = torch.reshape( |
| self.attention_cds(torch.concat([regFeatures, rnaCount], dim=1)), |
| [-1, self.hyperParam["filters1"], 1, self.motif_depth], |
| ) |
| cds_output = self.motif_detection(cdsSequence, attention_cds, "cds") |
| cds_seq_features = self.encoderDict["cds"](cds_output) |
| |
| TE = self.RPF_fc(cds_seq_features) |
| RPF = TE * rnaCount + self.bias |
| |
|
|
| return torch.concat([RPF, TE], 1) |
|
|
