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

Fix batch norm for conv2d

Browse source 
* Change tensorboard folders
* Save root scripts in output code
This commit is contained in:
Corentin 2021-01-30 02:11:06 +09:00
commit 42ae4474dd
2 changed files with 10 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

4
layers.py
View file

@ -85,7 +85,7 @@ class Conv2d(Layer):
self.batch_norm = nn.BatchNorm2d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=not Layer.BATCH_NORM_TRAINING) if self.batch_norm else None
track_running_stats=Layer.BATCH_NORM_TRAINING) if self.batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.conv(input_data))
@ -118,7 +118,7 @@ class Deconv2d(Layer):
self.batch_norm = nn.BatchNorm2d(
out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=not Layer.BATCH_NORM_TRAINING) if self.batch_norm else None
track_running_stats=Layer.BATCH_NORM_TRAINING) if self.batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.deconv(input_data))