def test_mask_type():
MaskedConv2d(1, 3, 3, mask_type="A")
MaskedConv2d(1, 3, 3, mask_type="B")
with pytest.raises(ValueError):
MaskedConv2d(1, 3, 3, mask_type="C")
def test_masked_conv_A(self):
conv = MaskedConv2d(3, 3, 3, padding=1)
with pytest.raises(RuntimeError):
m = torch.jit.script(conv)
def __init__(self,N=128,M=192,K=5,**kwargs): #'cuda:0' or 'cpu'
super().__init__(entropy_bottleneck_channels=N, **kwargs)
# super(DSIC, self).__init__()
# self.entropy_bottleneck1 = CompressionModel(entropy_bottleneck_channels=N)
# self.entropy_bottleneck2 = CompressionModel(entropy_bottleneck_channels=N)
self.gaussian1 = GaussianMixtureConditional(K = K)
self.gaussian2 = GaussianMixtureConditional(K = K)
self.N = int(N)
self.M = int(M)
self.K = int(K)
#定义组件
self.encoder1 = Encoder1(N,M)
self.encoder2 = Encoder2(N,M)
self.decoder1 = Decoder1(N,M)
self.decoder2 = Decoder2(N,M)
# pic2 需要的组件
# #hyper
# self._h_a1 = encode_hyper(N=N,M=M)
# self._h_a2 = encode_hyper(N=N,M=M)
# self._h_s1 = gmm_hyper_y1(N=N,M=M,K=K)
# self._h_s2 = gmm_hyper_y2(N=N,M=M,K=K)
######################################################################
self.h_a1 = nn.Sequential(
conv(M, N, stride=1, kernel_size=3),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
)
self.h_s1 = nn.Sequential(
deconv(N, M, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
deconv(M, M * 3 // 2, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(M * 3 // 2, M * 2, stride=1, kernel_size=3),
)
self.entropy_parameters1 = nn.Sequential(
nn.Conv2d(M * 12 // 3, M * 10 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 10 // 3, M * 8 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 8 // 3, M * 6 // 3, 1),
)
self.context_prediction1 = MaskedConv2d(M,
2 * M,
kernel_size=5,
padding=2,
stride=1)
self.gaussian_conditional1 = GaussianConditional(None)
self.h_a2 = nn.Sequential(
conv(M, N, stride=1, kernel_size=3),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
)
self.h_s2 = nn.Sequential(
deconv(N, M, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
deconv(M, M * 3 // 2, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(M * 3 // 2, M * 2, stride=1, kernel_size=3),
)
self.entropy_parameters2 = nn.Sequential(
nn.Conv2d(M * 18 // 3, M * 10 // 3, 1), # (M * 12 // 3, M * 10 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 10 // 3, M * 8 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 8 // 3, M * 6 // 3, 1),
)
self.context_prediction2 = MaskedConv2d(M,
2 * M,
kernel_size=5,
padding=2,
stride=1)
self.gaussian_conditional2 = GaussianConditional(None)
def test_mask_type():
MaskedConv2d(1, 3, 3, mask_type='A')
MaskedConv2d(1, 3, 3, mask_type='B')
with pytest.raises(ValueError):
MaskedConv2d(1, 3, 3, mask_type='C')
def __init__(self, N=192, **kwargs):
super().__init__(N=N, M=N, **kwargs)
self.g_a = nn.Sequential(
ResidualBlockWithStride(3, N, stride=2),
ResidualBlock(N, N),
ResidualBlockWithStride(N, N, stride=2),
ResidualBlock(N, N),
ResidualBlockWithStride(N, N, stride=2),
ResidualBlock(N, N),
conv3x3(N, N, stride=2),
)
self.g_a_b = copy.deepcopy(self.g_a)
self.h_a = nn.Sequential(
conv3x3(N, N),
nn.LeakyReLU(inplace=True),
conv3x3(N, N),
nn.LeakyReLU(inplace=True),
conv3x3(N, N, stride=2),
nn.LeakyReLU(inplace=True),
conv3x3(N, N),
nn.LeakyReLU(inplace=True),
conv3x3(N, N, stride=2),
)
self.h_a_b = copy.deepcopy(self.h_a)
self.h_s = nn.Sequential(
conv3x3(N, N),
nn.LeakyReLU(inplace=True),
subpel_conv3x3(N, N, 2),
nn.LeakyReLU(inplace=True),
conv3x3(N, N * 3 // 2),
nn.LeakyReLU(inplace=True),
subpel_conv3x3(N * 3 // 2, N * 3 // 2, 2),
nn.LeakyReLU(inplace=True),
conv3x3(N * 3 // 2, N * 2),
)
self.h_s_b = copy.deepcopy(self.h_s)
self.g_s = nn.Sequential(
ResidualBlock(N, N),
ResidualBlockUpsample(N, N, 2),
ResidualBlock(N, N),
ResidualBlockUpsample(N, N, 2),
ResidualBlock(N, N),
ResidualBlockUpsample(N, N, 2),
ResidualBlock(N, N),
subpel_conv3x3(N, 3, 2),
)
self.g_s_b = copy.deepcopy(self.g_s)
self.entropy_parameters = nn.Sequential(
nn.Conv2d(N * 4, 640, 3, padding=1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(640, 640, 3, padding=1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(640, N * 9, 3, padding=1),
)
self.context_prediction = MaskedConv2d(N,
2 * N,
kernel_size=5,
padding=2,
stride=1)
self.Quantization = Quantization()