Commit e85bb786 authored by Verena Praher's avatar Verena Praher
Browse files

add new model architectures

parent 3982785c
import torch.nn as nn
from utils import *
from datasets import MelSpecDataset
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from sklearn.metrics import roc_auc_score
# TODO pr-auc
# TODO f1-score
class CNN(pl.LightningModule):
def __init__(self, num_class):
super(CNN, self).__init__()
# init bn
self.bn_init = nn.BatchNorm2d(1)
# layer 1
self.conv_1 = nn.Conv2d(1, 64, 3, padding=1)
self.bn_1 = nn.BatchNorm2d(64)
self.mp_1 = nn.MaxPool2d((2, 4))
# layer 2
self.conv_2 = nn.Conv2d(64, 128, 3, padding=1)
self.bn_2 = nn.BatchNorm2d(128)
self.mp_2 = nn.MaxPool2d((2, 4))
# layer 3
self.conv_3 = nn.Conv2d(128, 128, 3, padding=1)
self.bn_3 = nn.BatchNorm2d(128)
self.mp_3 = nn.MaxPool2d((2, 4))
# layer 4
self.conv_4 = nn.Conv2d(128, 128, 3, padding=1)
self.bn_4 = nn.BatchNorm2d(128)
self.mp_4 = nn.MaxPool2d((3, 5))
# layer 5
self.conv_5 = nn.Conv2d(128, 64, 3, padding=1)
self.bn_5 = nn.BatchNorm2d(64)
self.mp_5 = nn.MaxPool2d((4, 4))
# classifier
self.dense = nn.Linear(64, num_class)
self.dropout = nn.Dropout(0.5)
def forward(self, x):
x = x.unsqueeze(1)
# init bn
x = self.bn_init(x)
# layer 1
x = self.mp_1(nn.ELU()(self.bn_1(self.conv_1(x))))
# layer 2
x = self.mp_2(nn.ELU()(self.bn_2(self.conv_2(x))))
# layer 3
x = self.mp_3(nn.ELU()(self.bn_3(self.conv_3(x))))
# layer 4
x = self.mp_4(nn.ELU()(self.bn_4(self.conv_4(x))))
# layer 5
x = self.mp_5(nn.ELU()(self.bn_5(self.conv_5(x))))
# classifier
x = x.view(x.size(0), -1)
x = self.dropout(x)
logit = nn.Sigmoid()(self.dense(x))
return logit
def my_loss(self, y_hat, y):
return F.binary_cross_entropy(y_hat, y)
\ No newline at end of file
from utils import *
from datasets import MelSpecDataset
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from sklearn.metrics import roc_auc_score
# TODO pr-auc
# TODO f1-score
from torchvision.models import resnet18
class Network(pl.LightningModule):
def __init__(self, num_tags):
super(Network, self).__init__()
self.num_tags = num_tags
self.model = resnet18(False)
num_features = self.model.fc.in_features
self.model.fc = nn.Linear(num_features, self.num_tags) # overwriting fc layer
self.sig = nn.Sigmoid(self.num_tags)
def forward(self, x):
x = self.model(x)
x = self.sig(x)
return x
def my_loss(self, y_hat, y):
return F.binary_cross_entropy(y_hat, y)
\ No newline at end of file
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment