Source code for imaginaire.generators.gancraft_base

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import functools
import re

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from imaginaire.layers import Conv2dBlock, LinearBlock
from imaginaire.model_utils.gancraft.layers import AffineMod, ModLinear
import imaginaire.model_utils.gancraft.mc_utils as mc_utils
import imaginaire.model_utils.gancraft.voxlib as voxlib
from imaginaire.utils.distributed import master_only_print as print


[docs]class RenderMLP(nn.Module): r""" MLP with affine modulation.""" def __init__(self, in_channels, style_dim, viewdir_dim, mask_dim=680, out_channels_s=1, out_channels_c=3, hidden_channels=256, use_seg=True): super(RenderMLP, self).__init__() self.use_seg = use_seg if self.use_seg: self.fc_m_a = nn.Linear(mask_dim, hidden_channels, bias=False) self.fc_viewdir = None if viewdir_dim > 0: self.fc_viewdir = nn.Linear(viewdir_dim, hidden_channels, bias=False) self.fc_1 = nn.Linear(in_channels, hidden_channels) self.fc_2 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True) self.fc_3 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True) self.fc_4 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True) self.fc_sigma = nn.Linear(hidden_channels, out_channels_s) if viewdir_dim > 0: self.fc_5 = nn.Linear(hidden_channels, hidden_channels, bias=False) self.mod_5 = AffineMod(hidden_channels, style_dim, mod_bias=True) else: self.fc_5 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True) self.fc_6 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True) self.fc_out_c = nn.Linear(hidden_channels, out_channels_c) self.act = nn.LeakyReLU(negative_slope=0.2)
[docs] def forward(self, x, raydir, z, m): r""" Forward network Args: x (N x H x W x M x in_channels tensor): Projected features. raydir (N x H x W x 1 x viewdir_dim tensor): Ray directions. z (N x style_dim tensor): Style codes. m (N x H x W x M x mask_dim tensor): One-hot segmentation maps. """ b, h, w, n, _ = x.size() z = z[:, None, None, None, :] f = self.fc_1(x) if self.use_seg: f = f + self.fc_m_a(m) # Common MLP f = self.act(f) f = self.act(self.fc_2(f, z)) f = self.act(self.fc_3(f, z)) f = self.act(self.fc_4(f, z)) # Sigma MLP sigma = self.fc_sigma(f) # Color MLP if self.fc_viewdir is not None: f = self.fc_5(f) f = f + self.fc_viewdir(raydir) f = self.act(self.mod_5(f, z)) else: f = self.act(self.fc_5(f, z)) f = self.act(self.fc_6(f, z)) c = self.fc_out_c(f) return sigma, c
[docs]class StyleMLP(nn.Module): r"""MLP converting style code to intermediate style representation.""" def __init__(self, style_dim, out_dim, hidden_channels=256, leaky_relu=True, num_layers=5, normalize_input=True, output_act=True): super(StyleMLP, self).__init__() self.normalize_input = normalize_input self.output_act = output_act fc_layers = [] fc_layers.append(nn.Linear(style_dim, hidden_channels, bias=True)) for i in range(num_layers-1): fc_layers.append(nn.Linear(hidden_channels, hidden_channels, bias=True)) self.fc_layers = nn.ModuleList(fc_layers) self.fc_out = nn.Linear(hidden_channels, out_dim, bias=True) if leaky_relu: self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True) else: self.act = functools.partial(F.relu, inplace=True)
[docs] def forward(self, z): r""" Forward network Args: z (N x style_dim tensor): Style codes. """ if self.normalize_input: z = F.normalize(z, p=2, dim=-1) for fc_layer in self.fc_layers: z = self.act(fc_layer(z)) z = self.fc_out(z) if self.output_act: z = self.act(z) return z
[docs]class SKYMLP(nn.Module): r"""MLP converting ray directions to sky features.""" def __init__(self, in_channels, style_dim, out_channels_c=3, hidden_channels=256, leaky_relu=True): super(SKYMLP, self).__init__() self.fc_z_a = nn.Linear(style_dim, hidden_channels, bias=False) self.fc1 = nn.Linear(in_channels, hidden_channels) self.fc2 = nn.Linear(hidden_channels, hidden_channels) self.fc3 = nn.Linear(hidden_channels, hidden_channels) self.fc4 = nn.Linear(hidden_channels, hidden_channels) self.fc5 = nn.Linear(hidden_channels, hidden_channels) self.fc_out_c = nn.Linear(hidden_channels, out_channels_c) if leaky_relu: self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True) else: self.act = functools.partial(F.relu, inplace=True)
[docs] def forward(self, x, z): r"""Forward network Args: x (... x in_channels tensor): Ray direction embeddings. z (... x style_dim tensor): Style codes. """ z = self.fc_z_a(z) while z.dim() < x.dim(): z = z.unsqueeze(1) y = self.act(self.fc1(x) + z) y = self.act(self.fc2(y)) y = self.act(self.fc3(y)) y = self.act(self.fc4(y)) y = self.act(self.fc5(y)) c = self.fc_out_c(y) return c
[docs]class RenderCNN(nn.Module): r"""CNN converting intermediate feature map to final image.""" def __init__(self, in_channels, style_dim, hidden_channels=256, leaky_relu=True): super(RenderCNN, self).__init__() self.fc_z_cond = nn.Linear(style_dim, 2 * 2 * hidden_channels) self.conv1 = nn.Conv2d(in_channels, hidden_channels, 1, stride=1, padding=0) self.conv2a = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1) self.conv2b = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1, bias=False) self.conv3a = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1) self.conv3b = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1, bias=False) self.conv4a = nn.Conv2d(hidden_channels, hidden_channels, 1, stride=1, padding=0) self.conv4b = nn.Conv2d(hidden_channels, hidden_channels, 1, stride=1, padding=0) self.conv4 = nn.Conv2d(hidden_channels, 3, 1, stride=1, padding=0) if leaky_relu: self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True) else: self.act = functools.partial(F.relu, inplace=True)
[docs] def modulate(self, x, w, b): w = w[..., None, None] b = b[..., None, None] return x * (w+1) + b
[docs] def forward(self, x, z): r"""Forward network. Args: x (N x in_channels x H x W tensor): Intermediate feature map z (N x style_dim tensor): Style codes. """ z = self.fc_z_cond(z) adapt = torch.chunk(z, 2 * 2, dim=-1) y = self.act(self.conv1(x)) y = y + self.conv2b(self.act(self.conv2a(y))) y = self.act(self.modulate(y, adapt[0], adapt[1])) y = y + self.conv3b(self.act(self.conv3a(y))) y = self.act(self.modulate(y, adapt[2], adapt[3])) y = y + self.conv4b(self.act(self.conv4a(y))) y = self.act(y) y = self.conv4(y) return y
[docs]class StyleEncoder(nn.Module): r"""Style Encoder constructor. Args: style_enc_cfg (obj): Style encoder definition file. """ def __init__(self, style_enc_cfg): super(StyleEncoder, self).__init__() input_image_channels = style_enc_cfg.input_image_channels num_filters = style_enc_cfg.num_filters kernel_size = style_enc_cfg.kernel_size padding = int(np.ceil((kernel_size - 1.0) / 2)) style_dims = style_enc_cfg.style_dims weight_norm_type = style_enc_cfg.weight_norm_type self.no_vae = getattr(style_enc_cfg, 'no_vae', False) activation_norm_type = 'none' nonlinearity = 'leakyrelu' base_conv2d_block = \ functools.partial(Conv2dBlock, kernel_size=kernel_size, stride=2, padding=padding, weight_norm_type=weight_norm_type, activation_norm_type=activation_norm_type, # inplace_nonlinearity=True, nonlinearity=nonlinearity) self.layer1 = base_conv2d_block(input_image_channels, num_filters) self.layer2 = base_conv2d_block(num_filters * 1, num_filters * 2) self.layer3 = base_conv2d_block(num_filters * 2, num_filters * 4) self.layer4 = base_conv2d_block(num_filters * 4, num_filters * 8) self.layer5 = base_conv2d_block(num_filters * 8, num_filters * 8) self.layer6 = base_conv2d_block(num_filters * 8, num_filters * 8) self.fc_mu = LinearBlock(num_filters * 8 * 4 * 4, style_dims) if not self.no_vae: self.fc_var = LinearBlock(num_filters * 8 * 4 * 4, style_dims)
[docs] def forward(self, input_x): r"""SPADE Style Encoder forward. Args: input_x (N x 3 x H x W tensor): input images. Returns: mu (N x C tensor): Mean vectors. logvar (N x C tensor): Log-variance vectors. z (N x C tensor): Style code vectors. """ if input_x.size(2) != 256 or input_x.size(3) != 256: input_x = F.interpolate(input_x, size=(256, 256), mode='bilinear') x = self.layer1(input_x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.layer5(x) x = self.layer6(x) x = x.view(x.size(0), -1) mu = self.fc_mu(x) if not self.no_vae: logvar = self.fc_var(x) std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) z = eps.mul(std) + mu else: z = mu logvar = torch.zeros_like(mu) return mu, logvar, z
[docs]class Base3DGenerator(nn.Module): r"""Minecraft 3D generator constructor. Args: gen_cfg (obj): Generator definition part of the yaml config file. data_cfg (obj): Data definition part of the yaml config file. """ def __init__(self, gen_cfg, data_cfg): super(Base3DGenerator, self).__init__() print('Base3DGenerator initialization.') # ---------------------- Main Network ------------------------ # Exclude some of the features from positional encoding self.pe_no_pe_feat_dim = getattr(gen_cfg, 'pe_no_pe_feat_dim', 0) # blk_feat passes through PE input_dim = (gen_cfg.blk_feat_dim-self.pe_no_pe_feat_dim)*(gen_cfg.pe_lvl_feat*2) + self.pe_no_pe_feat_dim if (gen_cfg.pe_incl_orig_feat): input_dim += (gen_cfg.blk_feat_dim-self.pe_no_pe_feat_dim) print('[Base3DGenerator] Expected input dimensions: ', input_dim) self.input_dim = input_dim self.mlp_model_kwargs = gen_cfg.mlp_model_kwargs self.pe_lvl_localcoords = getattr(gen_cfg, 'pe_lvl_localcoords', 0) if self.pe_lvl_localcoords > 0: self.mlp_model_kwargs['poscode_dim'] = self.pe_lvl_localcoords * 2 * 3 # Set pe_lvl_raydir=0 and pe_incl_orig_raydir=False to disable view direction input input_dim_viewdir = 3*(gen_cfg.pe_lvl_raydir*2) if (gen_cfg.pe_incl_orig_raydir): input_dim_viewdir += 3 print('[Base3DGenerator] Expected viewdir input dimensions: ', input_dim_viewdir) self.input_dim_viewdir = input_dim_viewdir self.pe_params = [gen_cfg.pe_lvl_feat, gen_cfg.pe_incl_orig_feat, gen_cfg.pe_lvl_raydir, gen_cfg.pe_incl_orig_raydir] # Style input dimension style_dims = gen_cfg.style_dims self.style_dims = style_dims interm_style_dims = getattr(gen_cfg, 'interm_style_dims', style_dims) self.interm_style_dims = interm_style_dims # ---------------------- Style MLP -------------------------- self.style_net = globals()[gen_cfg.stylenet_model]( style_dims, interm_style_dims, **gen_cfg.stylenet_model_kwargs) # number of output channels for MLP (before blending) final_feat_dim = getattr(gen_cfg, 'final_feat_dim', 16) self.final_feat_dim = final_feat_dim # ----------------------- Sky Network ------------------------- sky_input_dim_base = 3 # Dedicated sky network input dimensions sky_input_dim = sky_input_dim_base*(gen_cfg.pe_lvl_raydir_sky*2) if (gen_cfg.pe_incl_orig_raydir_sky): sky_input_dim += sky_input_dim_base print('[Base3DGenerator] Expected sky input dimensions: ', sky_input_dim) self.pe_params_sky = [gen_cfg.pe_lvl_raydir_sky, gen_cfg.pe_incl_orig_raydir_sky] self.sky_net = SKYMLP(sky_input_dim, style_dim=interm_style_dims, out_channels_c=final_feat_dim) # ----------------------- Style Encoder ------------------------- style_enc_cfg = getattr(gen_cfg, 'style_enc', None) setattr(style_enc_cfg, 'input_image_channels', 3) setattr(style_enc_cfg, 'style_dims', gen_cfg.style_dims) self.style_encoder = StyleEncoder(style_enc_cfg) # ---------------------- Ray Caster ------------------------- self.num_blocks_early_stop = gen_cfg.num_blocks_early_stop self.num_samples = gen_cfg.num_samples self.sample_depth = gen_cfg.sample_depth self.coarse_deterministic_sampling = getattr(gen_cfg, 'coarse_deterministic_sampling', True) self.sample_use_box_boundaries = getattr(gen_cfg, 'sample_use_box_boundaries', True) # ---------------------- Blender ------------------------- self.raw_noise_std = getattr(gen_cfg, 'raw_noise_std', 0.0) self.dists_scale = getattr(gen_cfg, 'dists_scale', 0.25) self.clip_feat_map = getattr(gen_cfg, 'clip_feat_map', True) self.keep_sky_out = getattr(gen_cfg, 'keep_sky_out', False) self.keep_sky_out_avgpool = getattr(gen_cfg, 'keep_sky_out_avgpool', False) keep_sky_out_learnbg = getattr(gen_cfg, 'keep_sky_out_learnbg', False) self.sky_global_avgpool = getattr(gen_cfg, 'sky_global_avgpool', False) if self.keep_sky_out: self.sky_replace_color = None if keep_sky_out_learnbg: sky_replace_color = torch.zeros([final_feat_dim]) sky_replace_color.requires_grad = True self.sky_replace_color = torch.nn.Parameter(sky_replace_color) # ---------------------- render_cnn ------------------------- self.denoiser = RenderCNN(final_feat_dim, style_dim=interm_style_dims) self.pad = gen_cfg.pad
[docs] def get_param_groups(self, cfg_opt): print('[Generator] get_param_groups') if hasattr(cfg_opt, 'ignore_parameters'): print('[Generator::get_param_groups] [x]: ignored.') optimize_parameters = [] for k, x in self.named_parameters(): match = False for m in cfg_opt.ignore_parameters: if re.match(m, k) is not None: match = True print(' [x]', k) break if match is False: print(' [v]', k) optimize_parameters.append(x) else: optimize_parameters = self.parameters() param_groups = [] param_groups.append({'params': optimize_parameters}) if hasattr(cfg_opt, 'param_groups'): optimized_param_names = [] all_param_names = [k for k, v in self.named_parameters()] param_groups = [] for k, v in cfg_opt.param_groups.items(): print('[Generator::get_param_groups] Adding param group from config:', k, v) params = getattr(self, k) named_parameters = [k] if issubclass(type(params), nn.Module): named_parameters = [k+'.'+pname for pname, _ in params.named_parameters()] params = params.parameters() param_groups.append({'params': params, **v}) optimized_param_names.extend(named_parameters) print('[Generator::get_param_groups] UNOPTIMIZED PARAMETERS:\n ', set(all_param_names) - set(optimized_param_names)) return param_groups
def _forward_perpix_sub(self, blk_feats, worldcoord2, raydirs_in, z, mc_masks_onehot=None): r"""Forwarding the MLP. Args: blk_feats (K x C1 tensor): Sparse block features. worldcoord2 (N x H x W x L x 3 tensor): 3D world coordinates of sampled points. raydirs_in (N x H x W x 1 x C2 tensor or None): ray direction embeddings. z (N x C3 tensor): Intermediate style vectors. mc_masks_onehot (N x H x W x L x C4): One-hot segmentation maps. Returns: net_out_s (N x H x W x L x 1 tensor): Opacities. net_out_c (N x H x W x L x C5 tensor): Color embeddings. """ proj_feature = voxlib.sparse_trilinear_interp_worldcoord( blk_feats, self.voxel.corner_t, worldcoord2, ign_zero=True) render_net_extra_kwargs = {} if self.pe_lvl_localcoords > 0: local_coords = torch.remainder(worldcoord2, 1.0) * 2.0 # Scale to [0, 2], as the positional encoding function doesn't have internal x2 local_coords[torch.isnan(local_coords)] = 0.0 local_coords = local_coords.contiguous() poscode = voxlib.positional_encoding(local_coords, self.pe_lvl_localcoords, -1, False) render_net_extra_kwargs['poscode'] = poscode if self.pe_params[0] == 0 and self.pe_params[1] is True: # no PE shortcut, saves ~400MB feature_in = proj_feature else: if self.pe_no_pe_feat_dim > 0: feature_in = voxlib.positional_encoding( proj_feature[..., :-self.pe_no_pe_feat_dim].contiguous(), self.pe_params[0], -1, self.pe_params[1]) feature_in = torch.cat([feature_in, proj_feature[..., -self.pe_no_pe_feat_dim:]], dim=-1) else: feature_in = voxlib.positional_encoding( proj_feature.contiguous(), self.pe_params[0], -1, self.pe_params[1]) net_out_s, net_out_c = self.render_net(feature_in, raydirs_in, z, mc_masks_onehot, **render_net_extra_kwargs) if self.raw_noise_std > 0.: noise = torch.randn_like(net_out_s) * self.raw_noise_std net_out_s = net_out_s + noise return net_out_s, net_out_c def _forward_perpix(self, blk_feats, voxel_id, depth2, raydirs, cam_ori_t, z): r"""Sample points along rays, forwarding the per-point MLP and aggregate pixel features Args: blk_feats (K x C1 tensor): Sparse block features. voxel_id (N x H x W x M x 1 tensor): Voxel ids from ray-voxel intersection test. M: num intersected voxels depth2 (N x 2 x H x W x M x 1 tensor): Depths of entrance and exit points for each ray-voxel intersection. raydirs (N x H x W x 1 x 3 tensor): The direction of each ray. cam_ori_t (N x 3 tensor): Camera origins. z (N x C3 tensor): Intermediate style vectors. """ # Generate sky_mask; PE transform on ray direction. with torch.no_grad(): raydirs_in = raydirs.expand(-1, -1, -1, 1, -1).contiguous() if self.pe_params[2] == 0 and self.pe_params[3] is True: raydirs_in = raydirs_in elif self.pe_params[2] == 0 and self.pe_params[3] is False: # Not using raydir at all raydirs_in = None else: raydirs_in = voxlib.positional_encoding(raydirs_in, self.pe_params[2], -1, self.pe_params[3]) # sky_mask: when True, ray finally hits sky sky_mask = voxel_id[:, :, :, [-1], :] == 0 # sky_only_mask: when True, ray hits nothing but sky sky_only_mask = voxel_id[:, :, :, [0], :] == 0 with torch.no_grad(): # Random sample points along the ray num_samples = self.num_samples + 1 if self.sample_use_box_boundaries: num_samples = self.num_samples - self.num_blocks_early_stop # 10 samples per ray + 4 intersections - 2 rand_depth, new_dists, new_idx = mc_utils.sample_depth_batched( depth2, num_samples, deterministic=self.coarse_deterministic_sampling, use_box_boundaries=self.sample_use_box_boundaries, sample_depth=self.sample_depth) worldcoord2 = raydirs * rand_depth + cam_ori_t[:, None, None, None, :] # Generate per-sample segmentation label voxel_id_reduced = self.label_trans.mc2reduced(voxel_id, ign2dirt=True) mc_masks = torch.gather(voxel_id_reduced, -2, new_idx) # B 256 256 N 1 mc_masks = mc_masks.long() mc_masks_onehot = torch.zeros([mc_masks.size(0), mc_masks.size(1), mc_masks.size( 2), mc_masks.size(3), self.num_reduced_labels], dtype=torch.float, device=voxel_id.device) # mc_masks_onehot: [B H W Nlayer 680] mc_masks_onehot.scatter_(-1, mc_masks, 1.0) net_out_s, net_out_c = self._forward_perpix_sub(blk_feats, worldcoord2, raydirs_in, z, mc_masks_onehot) # Handle sky sky_raydirs_in = raydirs.expand(-1, -1, -1, 1, -1).contiguous() sky_raydirs_in = voxlib.positional_encoding(sky_raydirs_in, self.pe_params_sky[0], -1, self.pe_params_sky[1]) skynet_out_c = self.sky_net(sky_raydirs_in, z) # Blending weights = mc_utils.volum_rendering_relu(net_out_s, new_dists * self.dists_scale, dim=-2) # If a ray exclusively hits the sky (no intersection with the voxels), set its weight to zero. weights = weights * torch.logical_not(sky_only_mask).float() total_weights_raw = torch.sum(weights, dim=-2, keepdim=True) # 256 256 1 1 total_weights = total_weights_raw is_gnd = worldcoord2[..., [0]] <= 1.0 # Y X Z, [256, 256, 4, 3], nan < 1.0 == False is_gnd = is_gnd.any(dim=-2, keepdim=True) nosky_mask = torch.logical_or(torch.logical_not(sky_mask), is_gnd) nosky_mask = nosky_mask.float() # Avoid sky leakage sky_weight = 1.0-total_weights if self.keep_sky_out: # keep_sky_out_avgpool overrides sky_replace_color if self.sky_replace_color is None or self.keep_sky_out_avgpool: if self.keep_sky_out_avgpool: if hasattr(self, 'sky_avg'): sky_avg = self.sky_avg else: if self.sky_global_avgpool: sky_avg = torch.mean(skynet_out_c, dim=[1, 2], keepdim=True) else: skynet_out_c_nchw = skynet_out_c.permute(0, 4, 1, 2, 3).squeeze(-1) sky_avg = F.avg_pool2d(skynet_out_c_nchw, 31, stride=1, padding=15, count_include_pad=False) sky_avg = sky_avg.permute(0, 2, 3, 1).unsqueeze(-2) # print(sky_avg.shape) skynet_out_c = skynet_out_c * (1.0-nosky_mask) + sky_avg*(nosky_mask) else: sky_weight = sky_weight * (1.0-nosky_mask) else: skynet_out_c = skynet_out_c * (1.0-nosky_mask) + self.sky_replace_color*(nosky_mask) if self.clip_feat_map is True: # intermediate feature before blending & CNN rgbs = torch.clamp(net_out_c, -1, 1) + 1 rgbs_sky = torch.clamp(skynet_out_c, -1, 1) + 1 net_out = torch.sum(weights*rgbs, dim=-2, keepdim=True) + sky_weight * \ rgbs_sky # 576, 768, 4, 3 -> 576, 768, 3 net_out = net_out.squeeze(-2) net_out = net_out - 1 elif self.clip_feat_map is False: rgbs = net_out_c rgbs_sky = skynet_out_c net_out = torch.sum(weights*rgbs, dim=-2, keepdim=True) + sky_weight * \ rgbs_sky # 576, 768, 4, 3 -> 576, 768, 3 net_out = net_out.squeeze(-2) elif self.clip_feat_map == 'tanh': rgbs = torch.tanh(net_out_c) rgbs_sky = torch.tanh(skynet_out_c) net_out = torch.sum(weights*rgbs, dim=-2, keepdim=True) + sky_weight * \ rgbs_sky # 576, 768, 4, 3 -> 576, 768, 3 net_out = net_out.squeeze(-2) else: raise NotImplementedError return net_out, new_dists, weights, total_weights_raw, rand_depth, net_out_s, net_out_c, skynet_out_c, \ nosky_mask, sky_mask, sky_only_mask, new_idx def _forward_global(self, net_out, z): r"""Forward the CNN Args: net_out (N x C5 x H x W tensor): Intermediate feature maps. z (N x C3 tensor): Intermediate style vectors. Returns: fake_images (N x 3 x H x W tensor): Output image. fake_images_raw (N x 3 x H x W tensor): Output image before TanH. """ fake_images = net_out.permute(0, 3, 1, 2) fake_images_raw = self.denoiser(fake_images, z) fake_images = torch.tanh(fake_images_raw) return fake_images, fake_images_raw