55 lines
2.1 KiB
Python
55 lines
2.1 KiB
Python
"""Universal network struture unit definition."""
|
|
from torch import nn
|
|
|
|
|
|
def define_squeeze_unit(basic_channel_size):
|
|
"""Define a 1x1 squeeze convolution with norm and activation."""
|
|
conv = nn.Conv2d(2 * basic_channel_size, basic_channel_size, kernel_size=1,
|
|
stride=1, padding=0, bias=False)
|
|
norm = nn.BatchNorm2d(basic_channel_size)
|
|
relu = nn.LeakyReLU(0.1)
|
|
layers = [conv, norm, relu]
|
|
return layers
|
|
|
|
|
|
def define_expand_unit(basic_channel_size):
|
|
"""Define a 3x3 expand convolution with norm and activation."""
|
|
conv = nn.Conv2d(basic_channel_size, 2 * basic_channel_size, kernel_size=3,
|
|
stride=1, padding=1, bias=False)
|
|
norm = nn.BatchNorm2d(2 * basic_channel_size)
|
|
relu = nn.LeakyReLU(0.1)
|
|
layers = [conv, norm, relu]
|
|
return layers
|
|
|
|
|
|
def define_halve_unit(basic_channel_size):
|
|
"""Define a 4x4 stride 2 expand convolution with norm and activation."""
|
|
conv = nn.Conv2d(basic_channel_size, 2 * basic_channel_size, kernel_size=4,
|
|
stride=2, padding=1, bias=False)
|
|
norm = nn.BatchNorm2d(2 * basic_channel_size)
|
|
relu = nn.LeakyReLU(0.1)
|
|
layers = [conv, norm, relu]
|
|
return layers
|
|
|
|
|
|
def define_depthwise_expand_unit(basic_channel_size):
|
|
"""Define a 3x3 expand convolution with norm and activation."""
|
|
conv1 = nn.Conv2d(basic_channel_size, 2 * basic_channel_size,
|
|
kernel_size=1, stride=1, padding=0, bias=False)
|
|
norm1 = nn.BatchNorm2d(2 * basic_channel_size)
|
|
relu1 = nn.LeakyReLU(0.1)
|
|
conv2 = nn.Conv2d(2 * basic_channel_size, 2 * basic_channel_size, kernel_size=3,
|
|
stride=1, padding=1, bias=False, groups=2 * basic_channel_size)
|
|
norm2 = nn.BatchNorm2d(2 * basic_channel_size)
|
|
relu2 = nn.LeakyReLU(0.1)
|
|
layers = [conv1, norm1, relu1, conv2, norm2, relu2]
|
|
return layers
|
|
|
|
|
|
def define_detector_block(basic_channel_size):
|
|
"""Define a unit composite of a squeeze and expand unit."""
|
|
layers = []
|
|
layers += define_squeeze_unit(basic_channel_size)
|
|
layers += define_expand_unit(basic_channel_size)
|
|
return layers
|