model.py

from torch import nn, cat
from torch.nn import init


# Create Network architecture
class MyNet(nn.Module):
    def __init__(self, x_shape, config):
        self.x_shape = x_shape
        super(MyNet, self).__init__()

        #  first cnn layers for amino

        self.l1 = nn.ModuleList()

        self.l1.append(
            nn.Conv1d(in_channels=1, out_channels=config['cnn_channels'], kernel_size=(config['kernel_size'],),
                      stride=(1,), padding=(int((config['kernel_size'] - 1) / 2),)))
        self.l1.append(nn.ELU())

        for layer in range(config['cnn_layers']):
            self.l1.append(nn.Conv1d(in_channels=config['cnn_channels'], out_channels=config['cnn_channels'],
                                     kernel_size=(config['kernel_size'],), stride=(1,),
                                     padding=(int((config['kernel_size'] - 1) / 2),)))
            self.l1.append(nn.ELU())
            self.l1.append(nn.Dropout(p=config['dropout'], inplace=False))

        self.l1.append(nn.Flatten())

        #  second cnn layers for diamino

        self.l2 = nn.ModuleList()

        self.l2.append(
            nn.Conv1d(in_channels=7, out_channels=config['cnn2_channels'], kernel_size=(config['kernel2_size'],),
                      stride=(1,), padding=(int((config['kernel2_size'] - 1) / 2),)))
        self.l2.append(nn.ELU())

        for layer in range(config['cnn2_layers']):
            self.l2.append(nn.Conv1d(in_channels=config['cnn2_channels'], out_channels=config['cnn2_channels'],
                                     kernel_size=(config['kernel2_size'],), stride=(1,),
                                     padding=(int((config['kernel2_size'] - 1) / 2),)))
            self.l2.append(nn.ELU())
            self.l2.append(nn.Dropout(p=config['dropout'], inplace=False))

        self.l2.append(nn.Flatten())

        #  third cnn layers for atoms

        self.l3 = nn.ModuleList()

        self.l3.append(
            nn.Conv1d(in_channels=6, out_channels=config['cnn3_channels'], kernel_size=(config['kernel3_size'],),
                      stride=(1,), padding=(int((config['kernel3_size'] - 1) / 2),)))
        self.l3.append(nn.ELU())

        for layer in range(config['cnn3_layers']):
            self.l3.append(nn.Conv1d(in_channels=config['cnn3_channels'], out_channels=config['cnn3_channels'],
                                     kernel_size=(config['kernel3_size'],), stride=(1,),
                                     padding=(int((config['kernel3_size'] - 1) / 2),)))
            self.l3.append(nn.ELU())
            self.l3.append(nn.Dropout(p=config['dropout'], inplace=False))

        self.l3.append(nn.Flatten())

        #  first fc layer for generale features

        self.l4 = nn.ModuleList()

        self.l4.append(nn.Linear(in_features=self.x_shape[2], out_features=config['fc_out']))
        self.l4.append(nn.ReLU())

        for layer in range(config['fc_layers']):
            self.l4.append(nn.Linear(in_features=config['fc_out'], out_features=config['fc_out']))
            self.l4.append(nn.ReLU())
        self.l4.append(nn.Flatten())

        #  fourth cnn layer for one hot encoder

        self.l5 = nn.ModuleList()

        self.l5.append(
            nn.Conv1d(in_channels=20, out_channels=config['cnn4_channels'], kernel_size=(config['kernel4_size'],),
                      stride=(1,), padding=(int((config['kernel4_size'] - 1) / 2),)))
        self.l5.append(nn.ELU())
        self.l5.append(nn.Dropout(p=config['dropout'], inplace=False))

        for layer in range(config['cnn4_layers']):
            self.l5.append(nn.Conv1d(in_channels=config['cnn4_channels'], out_channels=config['cnn4_channels'],
                                     kernel_size=(config['kernel4_size'],), stride=(1,),
                                     padding=(int((config['kernel4_size'] - 1) / 2),)))
            self.l5.append(nn.ELU())
            self.l5.append(nn.Dropout(p=config['dropout'], inplace=False))

        self.l5.append(nn.Flatten())

        #  final FC layer for concatenating all layers

        self.l6 = nn.ModuleList()

        self.l6.append(nn.Linear(in_features=(
                ((config['cnn_channels']) + config['cnn3_channels'] + config['cnn4_channels']) * self.x_shape[2] +
                config['fc_out'] * 7 + config['cnn2_channels'] * int(self.x_shape[2] / 2)),
            out_features=int(config['fc2_out'])))
        self.l6.append(nn.ELU())

        for layer in range(config['fc2_layers']):
            self.l6.append(nn.Linear(in_features=config['fc2_out'], out_features=config['fc2_out']))
            self.l6.append(nn.ELU())

        self.l6.append(nn.Linear(in_features=config['fc2_out'], out_features=int(config['fc2_out'])))
        self.l6.append(nn.ReLU())
        self.l6.append(nn.Linear(in_features=int(config['fc2_out']), out_features=int(config['fc2_out'])))
        self.l6.append(nn.ReLU())
        self.l6.append(nn.Linear(in_features=int(config['fc2_out']), out_features=1))

        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                init.kaiming_uniform_(m.weight, nonlinearity='leaky_relu')
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.kaiming_uniform_(m.weight, nonlinearity='leaky_relu')
                init.constant_(m.bias, 0)

    def forward(self, x):
        x1 = x[:, :1, :]
        x2 = x[:, 1:8, :int(self.x_shape[2] / 2)]
        x3 = x[:, 8:14, :]
        x4 = x[:, 14:21, :]
        x5 = x[:, 21:, :]

        for layer in self.l1:
            x1 = layer(x1)

        for layer in self.l2:
            x2 = layer(x2)

        for layer in self.l3:
            x3 = layer(x3)

        for layer in self.l4:
            x4 = layer(x4)

        for layer in self.l5:
            x5 = layer(x5)

        x6 = cat((x1, x2, x3, x4, x5), 1)
        for layer in self.l6:
            x6 = layer(x6)

        return x6