project Refactor, CNN Classifier Basics
This commit is contained in:
+10
-16
@@ -1,11 +1,7 @@
|
||||
from abc import ABC
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
import pytorch_lightning as pl
|
||||
from lib.modules.utils import AutoPad, Interpolate
|
||||
|
||||
#
|
||||
@@ -26,12 +22,12 @@ class ConvModule(nn.Module):
|
||||
conv_filters=64, conv_kernel=5, conv_stride=1, conv_padding=0):
|
||||
super(ConvModule, self).__init__()
|
||||
|
||||
# Module Paramters
|
||||
# Module Parameters
|
||||
self.in_shape = in_shape
|
||||
in_channels, height, width = in_shape[0], in_shape[1], in_shape[2]
|
||||
self.activation = activation()
|
||||
|
||||
# Convolution Paramters
|
||||
# Convolution Parameters
|
||||
self.padding = conv_padding
|
||||
self.stride = conv_stride
|
||||
|
||||
@@ -44,7 +40,7 @@ class ConvModule(nn.Module):
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm(x) if self.norm else x
|
||||
x = self.norm(x)
|
||||
|
||||
tensor = self.conv(x)
|
||||
tensor = self.dropout(tensor)
|
||||
@@ -72,10 +68,10 @@ class DeConvModule(nn.Module):
|
||||
self.in_shape = in_shape
|
||||
self.conv_filters = conv_filters
|
||||
|
||||
self.autopad = AutoPad() if autopad else lambda x: x
|
||||
self.interpolation = Interpolate(scale_factor=interpolation_scale) if interpolation_scale else lambda x: x
|
||||
self.norm = nn.BatchNorm2d(in_channels, eps=1e-04, affine=False) if normalize else lambda x: x
|
||||
self.dropout = nn.Dropout2d(dropout) if dropout else lambda x: x
|
||||
self.autopad = AutoPad() if autopad else lambda x: x
|
||||
self.interpolation = Interpolate(scale_factor=interpolation_scale) if interpolation_scale else lambda x: x
|
||||
self.norm = nn.BatchNorm2d(in_channels, eps=1e-04, affine=False) if normalize else lambda x: x
|
||||
self.dropout = nn.Dropout2d(dropout) if dropout else lambda x: x
|
||||
self.de_conv = nn.ConvTranspose2d(in_channels, self.conv_filters, conv_kernel, bias=use_bias,
|
||||
padding=self.padding, stride=self.stride)
|
||||
|
||||
@@ -100,13 +96,13 @@ class ResidualModule(nn.Module):
|
||||
output = self(x)
|
||||
return output.shape[1:]
|
||||
|
||||
def __init__(self, in_shape, module_class, n, activation=None, **module_paramters):
|
||||
def __init__(self, in_shape, module_class, n, activation=None, **module_parameters):
|
||||
assert n >= 1
|
||||
super(ResidualModule, self).__init__()
|
||||
self.in_shape = in_shape
|
||||
module_paramters.update(in_shape=in_shape)
|
||||
module_parameters.update(in_shape=in_shape)
|
||||
self.activation = activation() if activation else lambda x: x
|
||||
self.residual_block = nn.ModuleList([module_class(**module_paramters) for _ in range(n)])
|
||||
self.residual_block = nn.ModuleList([module_class(**module_parameters) for _ in range(n)])
|
||||
assert self.in_shape == self.shape, f'The in_shape: {self.in_shape} - must match the out_shape: {self.shape}.'
|
||||
|
||||
def forward(self, x):
|
||||
@@ -143,5 +139,3 @@ class RecurrentModule(nn.Module):
|
||||
def forward(self, x):
|
||||
tensor = self.rnn(x)
|
||||
return tensor
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user