"""
File: qumphy/models/alexnet.py
Project: 22HLT01 QUMPHY
Contact: oskar.pfeffer@ptb.de
Gitlab: https://gitlab.com/qumphy
Description: Alexnet pytorch lightning model.
"""
import torch
from torch import nn
import qumphy
[docs]
class AlexNet1D(nn.Module):
"""General 1D AlexNet model implementation.
The input_size and output size can be chosen freely.
Minimum input_size is 67.
The input shape is (batch_size, 1, input_size) and the output shape is
(batch_size, output_size).
Activation function for the last layer can be chosen freely.
"""
def __init__(self, input_size, output_size, output_activation=nn.Identity()):
"""Initialize the 1D AlexNet model.
Parameters
----------
input_size : int
Size of the one-dimensional input signal.
output_size : int
Size of the model output.
output_activation : nn.Module
Activation function applied to the output layer.
"""
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.output_activation = output_activation
self.feature_extraction = nn.Sequential(
nn.Conv1d(1, 96, kernel_size=11, stride=4, padding=0),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=2),
nn.Conv1d(96, 256, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=2),
nn.Conv1d(256, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv1d(384, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv1d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=2),
nn.Flatten(1, -1),
)
feature_size = self._get_feature_size()
self.feature_lerner = nn.Sequential(
nn.Linear(feature_size, 4096),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(4096, output_size),
)
[docs]
def forward(self, x):
"""Run a forward pass through the model.
Parameters
----------
x : torch.Tensor
Input tensor of shape (batch_size, 1, input_size).
Returns
-------
torch.Tensor
Model output tensor of shape (batch_size, output_size).
"""
x = self.feature_extraction(x)
x = self.feature_lerner(x)
x = self.output_activation(x)
return x
def _get_feature_size(self):
"""Get the size of the flattened feature vector.
Returns
-------
int
Size of the flattened feature vector after the convolutional layers.
"""
return self.feature_extraction(torch.rand((1, 1, self.input_size))).data.shape[
1
]
[docs]
class AlexNet1D_MCD(nn.Module):
"""General 1D AlexNet model implementation with Monte Carlo dropout.
The input_size and output size can be chosen freely.
Minimum input_size is 67.
The input shape is (batch_size, 1, input_size) and the output shape is
(batch_size, output_size).
Activation function for the last layer can be chosen freely.
"""
def __init__(
self,
input_size,
output_size,
dropout_rate=0.05,
output_activation=nn.Identity(),
mcdropout=True,
):
"""Initialize the 1D AlexNet model with Monte Carlo dropout.
Parameters
----------
input_size : int
Size of the one-dimensional input signal.
output_size : int
Size of the model output.
dropout_rate : float
Dropout probability used in the Monte Carlo dropout layers.
output_activation : nn.Module
Activation function applied to the output layer.
mcdropout : bool
If True, enables Monte Carlo dropout behavior.
"""
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.output_activation = output_activation
self.dropout_rate = dropout_rate
self.mcdropout = mcdropout
self.feature_extraction = nn.Sequential(
nn.Conv1d(1, 96, kernel_size=11, stride=4, padding=0),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=2),
nn.Conv1d(96, 256, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
qumphy.models.utils.mcdropout.MCDropout(self.dropout_rate, self.mcdropout),
nn.MaxPool1d(kernel_size=3, stride=2),
nn.Conv1d(256, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
qumphy.models.utils.mcdropout.MCDropout(self.dropout_rate, self.mcdropout),
nn.Conv1d(384, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
qumphy.models.utils.mcdropout.MCDropout(self.dropout_rate, self.mcdropout),
nn.Conv1d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
qumphy.models.utils.mcdropout.MCDropout(self.dropout_rate, self.mcdropout),
nn.MaxPool1d(kernel_size=3, stride=2),
nn.Flatten(1, -1),
)
feature_size = self._get_feature_size()
self.feature_lerner = nn.Sequential(
nn.Linear(feature_size, 4096),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(4096, output_size),
)
[docs]
def forward(self, x):
"""Run a forward pass through the model.
Parameters
----------
x : torch.Tensor
Input tensor of shape (batch_size, 1, input_size).
Returns
-------
torch.Tensor
Model output tensor of shape (batch_size, output_size).
"""
x = self.feature_extraction(x)
x = self.feature_lerner(x)
x = self.output_activation(x)
return x
def _get_feature_size(self):
"""Get the size of the flattened feature vector.
Returns
-------
int
Size of the flattened feature vector after the convolutional layers.
"""
return self.feature_extraction(torch.rand((1, 1, self.input_size))).data.shape[
1
]