def __init__(self, in_channels=3, out_channels=5):
super(VGG3D, self).__init__()
self.conv1 = Sequential(
Conv3d(in_channels, 64, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
InstanceNorm3d(64, True), ReLU(True),
Conv3d(64, 64, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
InstanceNorm3d(64, True), ReLU(True),
Conv3d(64, 64, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
InstanceNorm3d(64, True), ReLU(True),
MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2)))
self.conv2 = Sequential(
Conv3d(64, 128, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
InstanceNorm3d(128, True), ReLU(True),
Conv3d(128, 128, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
InstanceNorm3d(128, True), ReLU(True),
Conv3d(128, 128, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
InstanceNorm3d(128, True), ReLU(True),
MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2)))
self.conv3 = Sequential(Conv3d(128, 256, kernel_size=3, padding=1),
InstanceNorm3d(256, True), ReLU(True),
Conv3d(256, 256, kernel_size=3, padding=1),
InstanceNorm3d(256, True), ReLU(True),
Conv3d(256, 256, kernel_size=3, padding=1),
InstanceNorm3d(256, True), ReLU(True),
MaxPool3d(2, stride=2))
self.avgpool = AdaptiveAvgPool3d(1, 1, 1)
self.fc = Linear(256, out_channels)
def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=False):
super(Module, self).__init__()
if momentum is None:
momentum = 0.1
self._instance_norm = InstanceNorm3d(num_features,
eps=eps,
momentum=momentum,
affine=True)
def __init__(self, inc, outc):
super(DownsampleBlock, self).__init__()
self.conv = Conv3d(inc,
outc,
kernel_size=3,
stride=2,
padding=1,
bias=False)
self.norm = InstanceNorm3d(outc)
def __init__(self, inc, outc):
super(UpsampleBlock, self).__init__()
self.up = Upsample(scale_factor=2, mode="trilinear")
self.conv = Conv3d(inc,
outc,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.norm = InstanceNorm3d(outc)
def __init__(self, inc, outc):
super(FirstStem, self).__init__()
filters = outc // 2
self.conv1 = Conv3d(inc,
filters,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.norm1 = InstanceNorm3d(filters)
self.conv2 = Conv3d(filters,
outc,
kernel_size=3,
stride=2,
padding=1,
bias=False)
self.norm2 = InstanceNorm3d(outc)
def __init__(self, inc, outc):
super(FinalStem, self).__init__()
filters = outc * 2
self.conv1 = Conv3d(inc,
filters,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.norm1 = InstanceNorm3d(filters)
self.upsample = Upsample(scale_factor=2, mode="trilinear")
self.conv2 = Conv3d(filters,
outc,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.norm2 = InstanceNorm3d(outc)
def __init__(self, inc, outc, kernel_size=3, mode="3d"):
super(ConvBlock, self).__init__()
if mode not in ["2d", "2D", "3d", "3D", "p3d", "P3D"]:
raise ValueError("Unknow mode for convolution")
#endif
pad_size = (kernel_size - 1) // 2
if mode in ["2d", "2D"]:
self.conv1 = Conv3d(inc,
outc,
kernel_size=(kernel_size, kernel_size, 1),
stride=1,
padding=(pad_size, pad_size, 0),
bias=False)
self.conv = self.conv2d
#endif
if mode in ["3d", "3D"]:
self.conv1 = Conv3d(inc,
outc,
kernel_size=kernel_size,
stride=1,
padding=pad_size,
bias=False)
self.conv = self.conv3d
#endif
if mode in ["p3d", "P3D"]:
self.conv1 = Conv3d(inc,
outc,
kernel_size=(kernel_size, kernel_size, 1),
stride=1,
padding=(pad_size, pad_size, 0),
bias=False)
self.conv2 = Conv3d(outc,
outc,
kernel_size=(1, 1, kernel_size),
stride=1,
padding=(0, 0, pad_size),
bias=False)
self.conv = self.convp3d
#endif
self.norm = InstanceNorm3d(outc)
def create_norm(num_features):
"""Creates a normalization layer.
Note:
The normalization is configured via :meth:`pytorch_layers.Config.norm`,
and :attr:`pytorch_layers.Config.norm_kwargs`, and the saptial dimension
is configured via :attr:`pytorch_layers.Config.dim`.
Args:
num_features (int): The number of input channels.
Returns:
torch.nn.Module: The created normalization layer.
"""
config = Config()
if config.norm_mode is NormMode.GROUP:
from torch.nn import GroupNorm
kwargs = config.norm_kwargs.copy()
num_groups = kwargs.pop('num_groups')
return GroupNorm(num_groups, num_features, **kwargs)
elif config.norm_mode is NormMode.NONE:
from torch.nn import Identity
return Identity()
if config.dim is Dim.ONE:
if config.norm_mode is NormMode.INSTANCE:
from torch.nn import InstanceNorm1d
return InstanceNorm1d(num_features, **config.norm_kwargs)
elif config.norm_mode is NormMode.BATCH:
from torch.nn import BatchNorm1d
return BatchNorm1d(num_features, **config.norm_kwargs)
elif config.dim is Dim.TWO:
if config.norm_mode is NormMode.INSTANCE:
from torch.nn import InstanceNorm2d
return InstanceNorm2d(num_features, **config.norm_kwargs)
elif config.norm_mode is NormMode.BATCH:
from torch.nn import BatchNorm2d
return BatchNorm2d(num_features, **config.norm_kwargs)
elif config.dim is Dim.THREE:
if config.norm_mode is NormMode.INSTANCE:
from torch.nn import InstanceNorm3d
return InstanceNorm3d(num_features, **config.norm_kwargs)
elif config.norm_mode is NormMode.BATCH:
from torch.nn import BatchNorm3d
return BatchNorm3d(num_features, **config.norm_kwargs)