kxngluca/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

Allow noise to be injected at the generator inputs for resgen

Browse Source
This commit is contained in:
James Betker 2020-05-12 16:26:29 -06:00
parent 06d18343f7
commit 5d1b4caabf
2 changed files with 14 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
codes/models/archs/ResGen_arch.py
View File

@ -59,14 +59,18 @@ class FixupBasicBlock(nn.Module):
class FixupResNet(nn.Module):
def __init__(self, block, layers, upscale_applications=2, num_filters=64):
def __init__(self, block, layers, upscale_applications=2, num_filters=64, inject_noise=False):
super(FixupResNet, self).__init__()
self.num_layers = sum(layers) + layers[-1] # The last layer is applied twice to achieve 4x upsampling.
self.inplanes = num_filters
self.upscale_applications = upscale_applications
self.inject_noise = inject_noise
# Part 1 - Process raw input image. Most denoising should appear here and this should be the most complicated
# part of the block.
self.conv1 = nn.Conv2d(3, num_filters, kernel_size=5, stride=1, padding=2,
input_planes = 3
if inject_noise:
input_planes = 4
self.conv1 = nn.Conv2d(input_planes, num_filters, kernel_size=5, stride=1, padding=2,
bias=False)
self.bias1 = nn.Parameter(torch.zeros(1))
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
@ -119,6 +123,9 @@ class FixupResNet(nn.Module):
return nn.Sequential(*layers)
def forward(self, x):
if self.inject_noise:
rand_feature = torch.randn((x.shape[0], 1) + x.shape[2:], device=x.device, dtype=x.dtype)
x = torch.cat([x, rand_feature], dim=1)
x = self.conv1(x)
x = self.lrelu(x + self.bias1)
x = self.layer1(x)
@ -161,6 +168,9 @@ class FixupResNetV2(FixupResNet):
return x
def forward(self, x):
if self.inject_noise:
rand_feature = torch.randn((x.shape[0], 1) + x.shape[2:], device=x.device, dtype=x.dtype)
x = torch.cat([x, rand_feature], dim=1)
x = self.conv1(x)
x = self.lrelu(x + self.bias1)
x = self.layer1(x)

3
codes/models/networks.py
View File

@ -38,7 +38,8 @@ def define_G(opt):
upscale_applications=opt_net['upscale_applications'], num_filters=opt_net['nf'])
elif which_model == 'ResGenV2':
netG = ResGen_arch.fixup_resnet34_v2(nb_denoiser=opt_net['nb_denoiser'], nb_upsampler=opt_net['nb_upsampler'],
upscale_applications=opt_net['upscale_applications'], num_filters=opt_net['nf'])
upscale_applications=opt_net['upscale_applications'], num_filters=opt_net['nf'],
inject_noise=opt_net['inject_noise'])
# image corruption
elif which_model == 'HighToLowResNet':