def forward(self, z, xs=None):
"""
:param z: shape. = (B, z_dim)
:return:
"""
B, D = z.shape
map_h, map_w = 2, 2
decoder_out = torch.zeros(B,
D,
map_h,
map_w,
device=z.device,
dtype=z.dtype)
kl_losses = []
for i in range(len(self.map_from_z)):
# (B, m_h, m_w, z_dim)
z_rep = z.unsqueeze(1).repeat(1, map_h * map_w,
1).reshape(z.shape[0], map_h, map_w,
z.shape[1])
# (B, m_h, m_w, 2)
grid = create_grid(map_h, map_w, z.device).unsqueeze(0).repeat(
z.shape[0], 1, 1, 1)
grid = self.pos_map(grid)
# (B, z_dim, m_h, m_w)
z_sample = self.map_from_z[i](torch.cat(
[z_rep, grid], dim=-1).permute(0, 3, 1, 2).contiguous())
z_sample = torch.cat([decoder_out, z_sample], dim=1)
decoder_out = self.decoder_residual_blocks[i](
self.decoder_blocks[i](z_sample))
if i == len(self.map_from_z) - 1:
break
mu, log_var = self.condition_z[i](decoder_out).squeeze(-1).squeeze(
-1).chunk(2, dim=-1)
if xs is not None:
delta_mu, delta_log_var = self.condition_xz[i](torch.cat([xs[i], decoder_out], dim=1)) \
.squeeze(-1).squeeze(-1).chunk(2, dim=-1)
kl_losses.append(kl_2(delta_mu, delta_log_var, mu, log_var))
mu = mu + delta_mu
log_var = log_var + delta_log_var
z = reparameterize(mu, torch.exp(0.5 * log_var))
map_h *= 2**(len(self.decoder_blocks[i].channels) - 1)
map_w *= 2**(len(self.decoder_blocks[i].channels) - 1)
x_hat = self.recon(decoder_out)
return x_hat, kl_losses
def forward(self, z, xs=None, mode="random", freeze_level=-1):
"""
:param z: shape. = (B, z_dim, map_h, map_w)
:return:
"""
B, D, map_h, map_w = z.shape
# The init h (hidden state), can be replace with learned param, but it didn't work much
decoder_out = torch.zeros(B,
D,
map_h,
map_w,
device=z.device,
dtype=z.dtype)
kl_losses = []
if freeze_level != -1 and len(self.zs) == 0:
self.zs.append(z)
for i in range(len(self.decoder_residual_blocks)):
z_sample = torch.cat([decoder_out, z], dim=1)
decoder_out = self.decoder_residual_blocks[i](
self.decoder_blocks[i](z_sample))
if i == len(self.decoder_residual_blocks) - 1:
break
mu, log_var = self.condition_z[i](decoder_out).chunk(2, dim=1)
if xs is not None:
delta_mu, delta_log_var = self.condition_xz[i](torch.cat([xs[i], decoder_out], dim=1)) \
.chunk(2, dim=1)
kl_losses.append(kl_2(delta_mu, delta_log_var, mu, log_var))
mu = mu + delta_mu
log_var = log_var + delta_log_var
if mode == "fix" and i < freeze_level:
if len(self.zs) < freeze_level + 1:
z = reparameterize(mu, 0)
self.zs.append(z)
else:
z = self.zs[i + 1]
elif mode == "fix":
z = reparameterize(mu,
0 if i == 0 else torch.exp(0.5 * log_var))
else:
z = reparameterize(mu, torch.exp(0.5 * log_var))
map_h *= 2**(len(self.decoder_blocks[i].channels) - 1)
map_w *= 2**(len(self.decoder_blocks[i].channels) - 1)
x_hat = torch.sigmoid(self.recon(decoder_out))
return x_hat, kl_losses
def forward(self, z, xs=None):
"""
:param z: shape. = (B, z_dim)
:return:
"""
B, D, map_h, map_w = z.shape
# The init h (hidden state), can be replace with learned param, but it didn't work much
decoder_out = torch.zeros(B,
D,
map_h,
map_w,
device=z.device,
dtype=z.dtype)
kl_losses = []
for i in range(len(self.map_from_z)):
z_sample = torch.cat([decoder_out, z], dim=1)
decoder_out = self.decoder_residual_blocks[i](
self.decoder_blocks[i](z_sample))
if i == len(self.map_from_z) - 1:
break
mu, log_var = self.condition_z[i](decoder_out).chunk(2, dim=1)
if xs is not None:
delta_mu, delta_log_var = self.condition_xz[i](torch.cat([xs[i], decoder_out], dim=1)) \
.chunk(2, dim=1)
kl_losses.append(kl_2(delta_mu, delta_log_var, mu, log_var))
mu = mu + delta_mu
log_var = log_var + delta_log_var
z = reparameterize(mu, torch.exp(0.5 * log_var))
map_h *= 2**(len(self.decoder_blocks[i].channels) - 1)
map_w *= 2**(len(self.decoder_blocks[i].channels) - 1)
x_hat = torch.sigmoid(self.recon(decoder_out))
return x_hat, kl_losses