CompressAI
CompressAI copied to clipboard
Entropy model of hyper-latents
Hi,
What should I do if I want to use zero-mean gaussian Entropy Parameters for hyper-latents(like in y
)?
To do that, Can I use the self.gaussian_conditional
function as it is?
For example, if I use ScaleHyperprior
model, the code below.
(It's all the same as the existing code, and I only added the scale_z
part.)
class ScaleHyperprior(CompressionModel):
def __init__(self, N, M, **kwargs):
super().__init__(entropy_bottleneck_channels=N, **kwargs)
self.g_a = nn.Sequential(
conv(3, N),
GDN(N),
conv(N, N),
GDN(N),
conv(N, N),
GDN(N),
conv(N, M),
)
self.g_s = nn.Sequential(
deconv(M, N),
GDN(N, inverse=True),
deconv(N, N),
GDN(N, inverse=True),
deconv(N, N),
GDN(N, inverse=True),
deconv(N, 3),
)
self.h_a = nn.Sequential(
conv(M, N, stride=1, kernel_size=3),
nn.ReLU(inplace=True),
conv(N, N),
nn.ReLU(inplace=True),
conv(N, N),
)
self.h_s = nn.Sequential(
deconv(N, N),
nn.ReLU(inplace=True),
deconv(N, N),
nn.ReLU(inplace=True),
conv(N, M, stride=1, kernel_size=3),
nn.ReLU(inplace=True),
)
self.gaussian_conditional = GaussianConditional(None)
self.N = int(N)
self.M = int(M)
@property
def downsampling_factor(self) -> int:
return 2 ** (4 + 2)
def forward(self, x):
y = self.g_a(x)
z = self.h_a(torch.abs(y))
# z_hat, z_likelihoods = self.entropy_bottleneck(z)
scale_z = torch.abs(torch.ones(z.shape, device='cuda'))
z_hat, z_likelihoods = self.gaussian_conditional(z, scale_z)
scales_hat = self.h_s(z_hat)
y_hat, y_likelihoods = self.gaussian_conditional(y, scales_hat)
x_hat = self.g_s(y_hat)
return {
"x_hat": x_hat,
"likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
}
If possible, we should revise def forward
, def compress
, etc..
I think if this is possible, we can use it in the same way as y
...
Is the above method valid? (If the above code is used for training, execution is possible.) Or I would appreciate it if you could let me know the link to refer to.
I'm sorry it's a stupid question that's hard to answer. Thanks.