def __init__(
self, in_channels, start_channels, fac=1, sn: int = 0, device=None
):
super().__init__()
specnorm = _sn_to_specnorm(sn)
self.step1 = specnorm(
nn.Conv3d(
in_channels=in_channels,
out_channels=_facify(start_channels, fac),
kernel_size=5,
stride=1,
padding=compute_same_padding(5, 1, 1),
)
)
self.step2 = nn.LeakyReLU()
self.step3 = specnorm(
nn.Conv3d(
in_channels=_facify(start_channels, fac),
out_channels=_facify(start_channels*2, fac),
kernel_size=3,
stride=2,
padding=compute_same_padding(3, 1, 1),
)
)
self.step4 = nn.GroupNorm(1, _facify(start_channels*2, fac))
self.step5 = nn.LeakyReLU()
self.step6 = specnorm(
nn.Conv3d(
in_channels=_facify(start_channels*2, fac),
out_channels=_facify(start_channels*2, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1),
)
)
self.step7 = nn.GroupNorm(1, _facify(start_channels*2, fac))
self.step8 = nn.LeakyReLU()
self.step9 = specnorm(
nn.Conv3d(
in_channels=_facify(start_channels*2, fac),
out_channels=_facify(start_channels, fac),
kernel_size=4,
stride=1,
padding=0,
)
)
self.step10 = nn.GroupNorm(1, _facify(start_channels, fac))
self.step11 = nn.LeakyReLU()
self.mbd = MinibatchDiscrimination(_facify(start_channels, fac), _facify(start_channels/2, fac), 50)
self.to_critic_value = specnorm(
nn.Linear(
in_features=_facify(start_channels, fac)+_facify(start_channels/2, fac), out_features=1
)
)
def __init__(self, in_channels, fac=1, sn: int = 0, device=None):
super().__init__()
specnorm = _sn_to_specnorm(sn)
self.step1 = specnorm(
nn.Conv3d(
in_channels=in_channels,
out_channels=_facify(128, fac),
kernel_size=5,
stride=1,
padding=compute_same_padding(5, 1, 1),
))
self.step2 = nn.LeakyReLU()
self.step4 = Residual3DConvBlock(
in_channels=_facify(128, fac),
n_filters=_facify(256, fac),
kernel_size=3,
stride=2,
sn=sn,
device=device,
)
self.step5 = Residual3DConvBlock(
in_channels=_facify(256, fac),
n_filters=_facify(256, fac),
kernel_size=3,
stride=2,
sn=sn,
device=device,
)
self.step6 = specnorm(
nn.Conv3d(
in_channels=_facify(256, fac),
out_channels=_facify(256, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1),
))
self.step7 = nn.GroupNorm(1, _facify(256, fac))
self.step8 = nn.LeakyReLU()
self.step9 = specnorm(
nn.Conv3d(
in_channels=_facify(256, fac),
out_channels=_facify(128, fac),
kernel_size=4,
stride=1,
padding=0,
))
self.step10 = nn.GroupNorm(1, _facify(128, fac))
self.step11 = nn.LeakyReLU()
self.to_critic_value = specnorm(
nn.Linear(in_features=_facify(128, fac), out_features=1))
def __init__(self, n_filters_in, n_filters, kernel_size, sn=0):
super(Residual3DDeconvBlock, self).__init__()
specnorm = _sn_to_specnorm(sn)
self.n_filters_in = n_filters_in
self.n_filters = n_filters
self.kernel_size = kernel_size
self.conv = specnorm(
nn.Conv3d(n_filters_in, n_filters, kernel_size=1, stride=1))
same_padding = compute_same_padding(self.kernel_size, 1, dilation=1)
block_blocks = [
specnorm(
nn.Conv3d(
n_filters,
n_filters,
kernel_size=self.kernel_size,
stride=1,
padding=same_padding,
)),
nn.GroupNorm(1, num_channels=n_filters),
nn.LeakyReLU(),
specnorm(
nn.Conv3d(
n_filters,
n_filters,
kernel_size=self.kernel_size,
stride=1,
padding=same_padding,
)),
nn.GroupNorm(1, num_channels=n_filters),
]
self.block = nn.Sequential(*tuple(block_blocks))
self.nonlin = nn.LeakyReLU()
def __init__(
self,
n_input,
n_output,
start_channels,
fac=1,
sn: int = 0,
device=None,
):
super().__init__()
specnorm = _sn_to_specnorm(sn)
conv_blocks = [
specnorm(
nn.Conv3d(
in_channels=n_input,
out_channels=_facify(start_channels, fac),
kernel_size=5,
stride=1,
padding=2,
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels, fac)),
nn.LeakyReLU(),
specnorm(
nn.Conv3d(
in_channels=start_channels,
out_channels=_facify(start_channels/2, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1)
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels/2, fac)),
nn.LeakyReLU(),
specnorm(nn.Conv3d(_facify(start_channels / 2, fac), n_output, kernel_size=1, stride=1)),
nn.Sigmoid(),
]
self.conv = nn.Sequential(*tuple(conv_blocks)).to(device=device)
def __init__(
self,
in_channels,
n_filters,
kernel_size,
stride,
trans=False,
sn: int = 0,
device=None,
):
super(Residual3DConvBlock, self).__init__()
specnorm = _sn_to_specnorm(sn)
self.n_filters = n_filters
self.kernel_size = kernel_size
self.stride = stride
self.trans = trans
if self.stride != 1:
self.downsample = nn.Sequential(
nn.MaxPool3d(kernel_size=2, stride=2),
specnorm(
nn.Conv3d(
in_channels,
out_channels=self.n_filters,
kernel_size=1,
stride=1,
padding=0,
)),
)
elif self.trans:
self.downsample = nn.Sequential(
specnorm(
nn.Conv3d(
in_channels,
out_channels=self.n_filters,
kernel_size=1,
stride=1,
padding=0,
)))
else:
self.downsample = nn.Identity()
self.downsample = self.downsample.to(device=device)
same_padding = compute_same_padding(self.kernel_size,
self.stride,
dilation=1)
block_elements = [
specnorm(
nn.Conv3d(
in_channels=in_channels,
out_channels=self.n_filters,
kernel_size=self.kernel_size,
stride=self.stride,
padding=same_padding,
)),
nn.GroupNorm(1, num_channels=self.n_filters),
nn.LeakyReLU(),
specnorm(
nn.Conv3d(
in_channels=self.n_filters,
out_channels=self.n_filters,
kernel_size=self.kernel_size,
stride=1,
padding=same_padding,
)),
nn.GroupNorm(1, num_channels=self.n_filters),
]
self.block = nn.Sequential(*tuple(block_elements)).to(device=device)
self.nonlin = nn.LeakyReLU()
def __init__(
self,
z_dim,
n_input,
n_output,
start_channels,
fac=1,
sn: int = 0,
device=None,
):
super().__init__()
specnorm = _sn_to_specnorm(sn)
embed_condition_blocks = [
specnorm(
nn.Conv3d(
in_channels=n_input,
out_channels=_facify(start_channels, fac),
kernel_size=5,
stride=1,
padding=2,
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels, fac)),
nn.LeakyReLU(),
specnorm(
nn.Conv3d(
in_channels=start_channels,
out_channels=_facify(start_channels, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1)
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels, fac)),
nn.LeakyReLU(),
]
self.embed_condition = nn.Sequential(*tuple(embed_condition_blocks)).to(device=device)
downsample_cond_block = [
specnorm(
nn.Conv3d(
in_channels=start_channels,
out_channels=_facify(start_channels*2, fac),
kernel_size=3,
stride=2,
padding=compute_same_padding(3, 2, 1)
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels*2, fac)),
nn.LeakyReLU(),
]
self.downsample_cond = nn.Sequential(*tuple(downsample_cond_block)).to(device=device)
self.embed_noise_label = EmbedNoise(z_dim, _facify(start_channels*2, fac), sn=sn)
combined_block = [
specnorm(
nn.Conv3d(
in_channels=start_channels*4,
out_channels=_facify(start_channels*2, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1),
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels*2, fac)),
nn.LeakyReLU(),
]
self.combined = nn.Sequential(*tuple(combined_block)).to(device=device)
deconv_block = [
specnorm(
nn.Conv3d(
in_channels=start_channels*2,
out_channels=_facify(start_channels, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1),
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels, fac)),
nn.LeakyReLU(),
]
self.deconv = nn.Sequential(*tuple(deconv_block)).to(device=device)
to_image_blocks = [
specnorm(
nn.Conv3d(
in_channels=start_channels*2,
out_channels=_facify(start_channels, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1),
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels, fac)),
nn.LeakyReLU(),
specnorm(
nn.Conv3d(
in_channels=start_channels,
out_channels=_facify(start_channels / 2, fac),
kernel_size=3,
stride=1,
padding=compute_same_padding(3, 1, 1),
)
),
nn.GroupNorm(1, num_channels=_facify(start_channels/2, fac)),
nn.LeakyReLU(),
specnorm(nn.Conv3d(_facify(start_channels/2, fac), n_output, kernel_size=1, stride=1)),
nn.Sigmoid(),
]
self.to_image = nn.Sequential(*tuple(to_image_blocks)).to(device=device)