budui/raycv
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

add ConvTranspose2d in Conv2d

Browse Source
This commit is contained in:
budui 2020-10-13 10:31:03 +08:00
parent a6ffab1445
commit 611901cbdf
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
model/base/module.py
View File

@ -52,21 +52,23 @@ class LinearBlock(nn.Module):
class Conv2dBlock(nn.Module): class Conv2dBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int, bias=None, def __init__(self, in_channels: int, out_channels: int, bias=None,
activation_type="ReLU", norm_type="NONE", activation_type="ReLU", norm_type="NONE", additional_norm_kwargs=None,
additional_norm_kwargs=None, pre_activation=False, **conv_kwargs): pre_activation=False, use_transpose_conv=False, **conv_kwargs):
super().__init__() super().__init__()
self.norm_type = norm_type self.norm_type = norm_type
self.activation_type = activation_type self.activation_type = activation_type
self.pre_activation = pre_activation self.pre_activation = pre_activation
conv = nn.ConvTranspose2d if use_transpose_conv else nn.Conv2d
if pre_activation: if pre_activation:
self.normalization = _normalization(norm_type, in_channels, additional_norm_kwargs) self.normalization = _normalization(norm_type, in_channels, additional_norm_kwargs)
self.activation = _activation(activation_type, inplace=False) self.activation = _activation(activation_type, inplace=False)
self.convolution = nn.Conv2d(in_channels, out_channels, **conv_kwargs) self.convolution = conv(in_channels, out_channels, **conv_kwargs)
else: else:
# if caller not set bias, set bias automatically. # if caller not set bias, set bias automatically.
conv_kwargs["bias"] = _use_bias_checker(norm_type) if bias is None else bias conv_kwargs["bias"] = _use_bias_checker(norm_type) if bias is None else bias
self.convolution = nn.Conv2d(in_channels, out_channels, **conv_kwargs) self.convolution = conv(in_channels, out_channels, **conv_kwargs)
self.normalization = _normalization(norm_type, out_channels, additional_norm_kwargs) self.normalization = _normalization(norm_type, out_channels, additional_norm_kwargs)
self.activation = _activation(activation_type) self.activation = _activation(activation_type)