Приложение 2

from torch import nn
import torch

import numpy as np
from tqdm import tqdm_notebook as tqdm
from torch import optim
from torch.optim import lr_scheduler

import pandas as pd
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as trs
import matplotlib.pyplot as plt

# Обычная комбинация слоев: Сверточный, Батч норм, Релу
def conv_bn(inp, oup, stride = 1):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, padding = 1, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
                        )

# Слой развёртки генома
class Flatten(nn.Module):
    def __init__(self, oup):
        self.oup = oup

    def __call__(self, sample):
        return sample.view(sample.shape[0], 1, self.oup, -1)


# Геномный слой по аналогии с basset
def genome_conv(previous_oup = 3, oup = 3):
    return nn.Sequential(
                                nn.Conv2d(1, oup, kernel_size = (previous_oup, 3)),
                                nn.ReLU(inplace = True),
                                Flatten(oup=oup),
                                nn.MaxPool2d(kernel_size= (1, 2))
                        )

# Если параметр mode = 'Genome' после первого слоя дальше идут "Геномные слои" genome_conv, если mode = Traditional, то conv_bn
class GenomeNet(nn.Module):
    def __init__(self, mode = 'Genome'):
        super(GenomeNet, self).__init__()

        self.mode = mode

        self.conv = nn.Conv2d(1, 3, kernel_size = (4,11))
        self.relu = nn.ReLU(inplace = True)
        self.pooling_initial = nn.MaxPool2d(kernel_size= (1, 4))
        self.pooling_afterwards = nn.MaxPool2d(kernel_size= (1, 2))


        if mode == 'Traditional':
        ### option 1
            self.ConvBnRelu1 = conv_bn(1, 32, stride = 1)
            self.ConvBnRelu2 = conv_bn(32, 64, stride = 1)
            self.ConvBnRelu3 = conv_bn(64, 128, stride = 1)
            self.ConvBnRelu4 = conv_bn(128, 260, stride = 1)
            self.GAP = nn.AdaptiveAvgPool2d(1)

        elif mode == 'Genome':
        ### option 2
            self.GenomeConv1 = genome_conv(previous_oup = 3, oup = 8)
            self.GenomeConv2 = genome_conv(previous_oup = 8, oup = 16)
            self.GenomeConv3 = genome_conv(previous_oup = 16, oup = 16)
            self.GenomeConv4 = genome_conv(previous_oup = 16, oup = 32)
            self.GenomeConv5 = genome_conv(previous_oup = 32, oup = 52)


        # Полносвязные слои
        self.FC1 = nn.Linear(260, 164)
        self.FC2 = nn.Linear(164, 164)

        self.FinalFC = nn.Linear(164, 1)
        self.Sigmoid = nn.Sigmoid()

    def forward(self, x):

        x = self.conv(x).reshape(x.shape[0], 1, 3, -1)

        x = self.relu(x)
        x = self.pooling_initial(x) #100 x 1 x 3 x 122

        if self.mode == 'Traditional':
            x = self.pooling_afterwards(self.ConvBnRelu1(x))
            x = self.pooling_afterwards(self.ConvBnRelu2(x))
            x = self.pooling_afterwards(self.ConvBnRelu3(x))
            x = self.pooling_afterwards(self.ConvBnRelu4(x))

            x = self.GAP(x).reshape(x.shape[0], -1)

        elif self.mode == 'Genome':
            x = self.GenomeConv1(x)
            x = self.GenomeConv2(x)
            x = self.GenomeConv3(x)
            x = self.GenomeConv4(x)
            x = self.GenomeConv5(x)

            x = x.reshape(x.shape[0], -1)

        x = self.relu(self.FC1(x))
        x = self.relu(self.FC2(x))

        x = self.Sigmoid(self.FinalFC(x))

        return x