def forward(self, x: torch.Tensor, y: torch.Tensor,
s_x: torch.Tensor) -> torch.Tensor:
assert utils.is_valid_image_tensor(x)
assert utils.is_valid_tensor(y, num_dims=2, batch_size=x.shape[0])
assert utils.is_valid_tensor(s_x, num_dims=2, batch_size=x.shape[0])
e_x = self._e(x, y)
l_sty = self._loss_fn(e_x, s_x)
return l_sty
def forward(self, z: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
assert utils.is_valid_tensor(z, num_dims=2)
assert utils.is_valid_tensor(y, num_dims=2, batch_size=z.shape[0])
x = z
for layer in self._body:
x = layer(x)
x = nn.functional.relu(x)
style_code = [layer(x) for layer in self._heads]
style_code = torch.stack(style_code, dim=1)
style_code = torch.cat([style_code[i, y_i] for i, y_i in enumerate(y)])
return style_code
def forward(
self, x_real: torch.Tensor, y_real: torch.Tensor, x_fake: torch.Tensor,
y_fake: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
assert utils.is_valid_image_tensor(x_real)
assert utils.is_valid_tensor(y_real,
num_dims=2,
batch_size=x_real.shape[0])
assert utils.is_valid_image_tensor(x_fake)
assert utils.is_valid_tensor(y_fake,
num_dims=2,
batch_size=x_fake.shape[0])
x_real.requires_grad = True
logits_real = self._d(x_real, y_real)
logits_fake = self._d(x_fake, y_fake)
l_adv_d = (
torch.mean(self._bce(logits_real, torch.ones_like(logits_real)),
dim=0) +
torch.mean(self._bce(logits_fake, torch.zeros_like(logits_fake)),
dim=0))
l_adv_g = torch.mean(self._bce(logits_fake,
torch.ones_like(logits_fake)),
dim=0)
# R1 regularization
grads_x_real = torch.autograd.grad(
outputs=logits_real.sum(),
inputs=x_real,
create_graph=True,
)[0]
r1_reg = 0.5 * grads_x_real.view(x_real.shape[0], -1).norm(
2, dim=1).mean()
return l_adv_d, l_adv_g, r1_reg
def forward(self, x: torch.Tensor,
style_code: torch.Tensor) -> torch.Tensor:
assert utils.is_valid_tensor(x, num_dims=4)
assert x.shape[1] == self._num_channels
x = self._instance_norm(x)
stats = self._linear(style_code)
mu = stats[:, :self._num_channels]
mu.unsqueeze_(2).unsqueeze_(2)
sigma = stats[:, self._num_channels:]
sigma.unsqueeze_(2).unsqueeze_(2)
x = x * sigma + mu
return x
def forward(self, x_real: torch.Tensor, y_real: torch.Tensor,
x_fake: torch.Tensor) -> torch.Tensor:
assert utils.is_valid_image_tensor(x_real)
assert utils.is_valid_image_tensor(x_fake)
assert x_real.shape[0] == x_fake.shape[0]
assert utils.is_valid_tensor(y_real,
num_dims=2,
batch_size=x_real.shape[0])
s_x = self._e(x_real, y_real)
x_recon = self._g(x_fake, s_x)
l_cyc = self._loss_fn(x_real, x_recon)
return l_cyc
def forward(self, x: torch.Tensor,
style_code: torch.Tensor) -> torch.Tensor:
assert utils.is_valid_image_tensor(x)
assert utils.is_valid_tensor(style_code,
num_dims=2,
batch_size=x.shape[0])
for block in self._all_blocks:
if isinstance(block, ResBlock):
x = block(x, style_code)
else:
x = block(x)
return x
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
assert utils.is_valid_image_tensor(x)
assert utils.is_valid_tensor(y, num_dims=2, batch_size=x.shape[0])
for block in self._down_blocks:
if isinstance(block, ResBlock):
x = block(x, None)
else:
x = block(x)
logits = [layer(x) for layer in self._heads]
outputs = torch.stack(logits, dim=1)
outputs = torch.cat([outputs[i, y_i] for i, y_i in enumerate(y)])
return outputs