Spaces:
Running
Running
File size: 5,923 Bytes
ce190ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import climategan.strings as strings
from climategan.blocks import InterpolateNearest2d, SPADEResnetBlock
from climategan.norms import SpectralNorm
def create_painter(opts, no_init=False, verbose=0):
if verbose > 0:
print(" - Add PainterSpadeDecoder Painter")
return PainterSpadeDecoder(opts)
class PainterSpadeDecoder(nn.Module):
def __init__(self, opts):
"""Create a SPADE-based decoder, which forwards z and the conditioning
tensors seg (in the original paper, conditioning is on a semantic map only).
All along, z is conditioned on seg. First 3 SpadeResblocks (SRB) do not shrink
the channel dimension, and an upsampling is applied after each. Therefore
2 upsamplings at this point. Then, for each remaining upsamplings
(w.r.t. spade_n_up), the SRB shrinks channels by 2. Before final conv to get 3
channels, the number of channels is therefore:
final_nc = channels(z) * 2 ** (spade_n_up - 2)
Args:
latent_dim (tuple): z's shape (only the number of channels matters)
cond_nc (int): conditioning tensor's expected number of channels
spade_n_up (int): Number of total upsamplings from z
spade_use_spectral_norm (bool): use spectral normalization?
spade_param_free_norm (str): norm to use before SPADE de-normalization
spade_kernel_size (int): SPADE conv layers' kernel size
Returns:
[type]: [description]
"""
super().__init__()
latent_dim = opts.gen.p.latent_dim
cond_nc = 3
spade_n_up = opts.gen.p.spade_n_up
spade_use_spectral_norm = opts.gen.p.spade_use_spectral_norm
spade_param_free_norm = opts.gen.p.spade_param_free_norm
spade_kernel_size = 3
self.z_nc = latent_dim
self.spade_n_up = spade_n_up
self.z_h = self.z_w = None
self.fc = nn.Conv2d(3, latent_dim, 3, padding=1)
self.head_0 = SPADEResnetBlock(
self.z_nc,
self.z_nc,
cond_nc,
spade_use_spectral_norm,
spade_param_free_norm,
spade_kernel_size,
)
self.G_middle_0 = SPADEResnetBlock(
self.z_nc,
self.z_nc,
cond_nc,
spade_use_spectral_norm,
spade_param_free_norm,
spade_kernel_size,
)
self.G_middle_1 = SPADEResnetBlock(
self.z_nc,
self.z_nc,
cond_nc,
spade_use_spectral_norm,
spade_param_free_norm,
spade_kernel_size,
)
self.up_spades = nn.Sequential(
*[
SPADEResnetBlock(
self.z_nc // 2 ** i,
self.z_nc // 2 ** (i + 1),
cond_nc,
spade_use_spectral_norm,
spade_param_free_norm,
spade_kernel_size,
)
for i in range(spade_n_up - 2)
]
)
self.final_nc = self.z_nc // 2 ** (spade_n_up - 2)
self.final_spade = SPADEResnetBlock(
self.final_nc,
self.final_nc,
cond_nc,
spade_use_spectral_norm,
spade_param_free_norm,
spade_kernel_size,
)
self.final_shortcut = None
if opts.gen.p.use_final_shortcut:
self.final_shortcut = nn.Sequential(
*[
SpectralNorm(nn.Conv2d(self.final_nc, 3, 1)),
nn.BatchNorm2d(3),
nn.LeakyReLU(0.2, True),
]
)
self.conv_img = nn.Conv2d(self.final_nc, 3, 3, padding=1)
self.upsample = InterpolateNearest2d(scale_factor=2)
def set_latent_shape(self, shape, is_input=True):
"""
Sets the latent shape to start the upsampling from, i.e. z_h and z_w.
If is_input is True, then this is the actual input shape which should
be divided by 2 ** spade_n_up
Otherwise, just sets z_h and z_w from shape[-2] and shape[-1]
Args:
shape (tuple): The shape to start sampling from.
is_input (bool, optional): Whether to divide shape by 2 ** spade_n_up
"""
if isinstance(shape, (list, tuple)):
self.z_h = shape[-2]
self.z_w = shape[-1]
elif isinstance(shape, int):
self.z_h = self.z_w = shape
else:
raise ValueError("Unknown shape type:", shape)
if is_input:
self.z_h = self.z_h // (2 ** self.spade_n_up)
self.z_w = self.z_w // (2 ** self.spade_n_up)
def _apply(self, fn):
# print("Applying SpadeDecoder", fn)
super()._apply(fn)
# self.head_0 = fn(self.head_0)
# self.G_middle_0 = fn(self.G_middle_0)
# self.G_middle_1 = fn(self.G_middle_1)
# for i, up in enumerate(self.up_spades):
# self.up_spades[i] = fn(up)
# self.conv_img = fn(self.conv_img)
return self
def forward(self, z, cond):
if z is None:
assert self.z_h is not None and self.z_w is not None
z = self.fc(F.interpolate(cond, size=(self.z_h, self.z_w)))
y = self.head_0(z, cond)
y = self.upsample(y)
y = self.G_middle_0(y, cond)
y = self.upsample(y)
y = self.G_middle_1(y, cond)
for i, up in enumerate(self.up_spades):
y = self.upsample(y)
y = up(y, cond)
if self.final_shortcut is not None:
cond = self.final_shortcut(y)
y = self.final_spade(y, cond)
y = self.conv_img(F.leaky_relu(y, 2e-1))
y = torch.tanh(y)
return y
def __str__(self):
return strings.spadedecoder(self)
|