model.py

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.autograd import Variable
from torch import optim
import numpy

class LSTMModel(nn.Module):

    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first = True)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)

        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])

        return out

class CNNModel(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size = (32, 2))
        self.pooling1 = nn.MaxPool2d(4)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size = (32, 1))
        self.pooling2 = nn.MaxPool2d(4)
        self.conv3 = nn.Conv2d(out_channels, out_channels, kernel_size = (32, 1))
        self.pooling3 = nn.MaxPool2d(4)
        self.conv4 = nn.Conv2d(out_channels, out_channels, kernel_size = (32, 1))
        self.pooling4 = nn.MaxPool2d(4)
        self.conv5 = nn.Conv2d(out_channels, out_channels, kernel_size = (32, 1))
        self.pooling5 = nn.MaxPool2d(4)
        self.fc = nn.Linear(270, 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.pooling1(x)
        x = self.conv2(x)
        x = self.pooling2(x)
        x = self.conv3(x)
        x = self.pooling2(x)
        x = self.conv4(x)
        x = self.pooling2(x)
        x = self.conv5(x)
        x = self.pooling2(x)
        x = x.view(-1, 270)
        x = self.fc(x)
        return torch.sigmoid(x)