corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Avoid use_batch_norm as layers instance variable

Browse source
This commit is contained in:
Corentin 2021-05-21 16:00:16 +09:00
commit 770a9a4f82
2 changed files with 57 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

31
layers.py
View file

@ -19,7 +19,7 @@ class Layer(nn.Module):
METRICS = False
LOGGER = DummyLogger()
def __init__(self, activation, use_batch_norm):
def __init__(self, activation):
super().__init__()
# Preload default
if activation == 0:
@ -28,28 +28,27 @@ class Layer(nn.Module):
self.activation = activation()
else:
self.activation = activation
self.batch_norm: torch.nn._BatchNorm = None
self.use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
self.batch_norm: torch.nn._BatchNorm
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
output = input_data
if self.activation is not None:
if self.activation:
output = self.activation(output)
if self.use_batch_norm:
# It is assumed here that if using batch norm, then self.batch_norm has been instanciated.
if self.batch_norm:
output = self.batch_norm(output)
return output
class Linear(Layer):
def __init__(self, in_channels: int, out_channels: int, activation=0, use_batch_norm: bool = None, **kwargs):
super().__init__(activation, use_batch_norm)
super().__init__(activation)
self.fc = nn.Linear(in_channels, out_channels, bias=not self.batch_norm, **kwargs)
use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
self.batch_norm = nn.BatchNorm1d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=Layer.BATCH_NORM_TRAINING) if self.use_batch_norm else None
track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.fc(input_data))
@ -58,14 +57,15 @@ class Linear(Layer):
class Conv1d(Layer):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
stride: Union[int, Tuple[int, int]] = 1, activation=0, use_batch_norm: bool = None, **kwargs):
super().__init__(activation, use_batch_norm)
super().__init__(activation)
self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride,
bias=not self.use_batch_norm, **kwargs)
use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
self.batch_norm = nn.BatchNorm1d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=Layer.BATCH_NORM_TRAINING) if self.use_batch_norm else None
track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.conv(input_data))
@ -78,6 +78,7 @@ class Conv2d(Layer):
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
bias=not self.use_batch_norm, **kwargs)
use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
self.batch_norm = nn.BatchNorm2d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
@ -90,14 +91,15 @@ class Conv2d(Layer):
class Conv3d(Layer):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
stride: Union[int, Tuple[int, int]] = 1, activation=0, use_batch_norm: bool = None, **kwargs):
super().__init__(activation, use_batch_norm)
super().__init__(activation)
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride,
bias=not self.use_batch_norm, **kwargs)
use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
self.batch_norm = nn.BatchNorm3d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=Layer.BATCH_NORM_TRAINING) if self.use_batch_norm else None
track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.conv(input_data))
@ -106,15 +108,16 @@ class Conv3d(Layer):
class Deconv2d(Layer):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
stride: Union[int, Tuple[int, int]] = 1, activation=0, use_batch_norm: bool = None, **kwargs):
super().__init__(activation, use_batch_norm)
super().__init__(activation)
self.deconv = nn.ConvTranspose2d(
in_channels, out_channels, kernel_size, stride=stride,
bias=not self.use_batch_norm, **kwargs)
use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
self.batch_norm = nn.BatchNorm2d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=Layer.BATCH_NORM_TRAINING) if self.use_batch_norm else None
track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.deconv(input_data))