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:
parent
343af70a8d
commit
f389025b53
|
@ -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)
|
||||||
|
|
Loading…
Reference in New Issue
Block a user