Change ResGen noise feature

It now injects noise directly into the input filters, rather than a
pure noise filter. The pure noise filter was producing really
poor results (and I'm honestly not quite sure why).
This commit is contained in:
James Betker 2020-05-13 09:22:06 -06:00
parent 343af70a8d
commit f389025b53

View File

@ -61,16 +61,13 @@ class FixupResNet(nn.Module):
def __init__(self, block, layers, upscale_applications=2, num_filters=64, inject_noise=False): def __init__(self, block, layers, upscale_applications=2, num_filters=64, inject_noise=False):
super(FixupResNet, self).__init__() super(FixupResNet, self).__init__()
self.inject_noise = inject_noise
self.num_layers = sum(layers) + layers[-1] # The last layer is applied twice to achieve 4x upsampling. self.num_layers = sum(layers) + layers[-1] # The last layer is applied twice to achieve 4x upsampling.
self.inplanes = num_filters self.inplanes = num_filters
self.upscale_applications = upscale_applications 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 1 - Process raw input image. Most denoising should appear here and this should be the most complicated
# part of the block. # part of the block.
input_planes = 3 self.conv1 = nn.Conv2d(3, num_filters, kernel_size=5, stride=1, padding=2,
if inject_noise:
input_planes = 4
self.conv1 = nn.Conv2d(input_planes, num_filters, kernel_size=5, stride=1, padding=2,
bias=False) bias=False)
self.bias1 = nn.Parameter(torch.zeros(1)) self.bias1 = nn.Parameter(torch.zeros(1))
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
@ -124,8 +121,8 @@ class FixupResNet(nn.Module):
def forward(self, x): def forward(self, x):
if self.inject_noise: if self.inject_noise:
rand_feature = torch.randn((x.shape[0], 1) + x.shape[2:], device=x.device, dtype=x.dtype) rand_feature = torch.randn_like(x)
x = torch.cat([x, rand_feature], dim=1) x = x + rand_feature * .1
x = self.conv1(x) x = self.conv1(x)
x = self.lrelu(x + self.bias1) x = self.lrelu(x + self.bias1)
x = self.layer1(x) x = self.layer1(x)