|
|
|
@ -46,7 +46,7 @@ class RRDB(nn.Module):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RRDBNet(nn.Module):
|
|
|
|
|
def __init__(self, in_nc, out_nc, nf, nb, gc=32):
|
|
|
|
|
def __init__(self, in_nc, out_nc, nf, nb, gc=32, interpolation_scale_factor=2):
|
|
|
|
|
super(RRDBNet, self).__init__()
|
|
|
|
|
RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
|
|
|
|
|
|
|
|
|
@ -61,13 +61,15 @@ class RRDBNet(nn.Module):
|
|
|
|
|
|
|
|
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
|
|
|
|
|
|
|
|
self.interpolation_scale_factor = interpolation_scale_factor
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
fea = self.conv_first(x)
|
|
|
|
|
trunk = self.trunk_conv(self.RRDB_trunk(fea))
|
|
|
|
|
fea = fea + trunk
|
|
|
|
|
|
|
|
|
|
fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
|
|
|
|
|
fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
|
|
|
|
|
fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=self.interpolation_scale_factor, mode='nearest')))
|
|
|
|
|
fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=self.interpolation_scale_factor, mode='nearest')))
|
|
|
|
|
out = self.conv_last(self.lrelu(self.HRconv(fea)))
|
|
|
|
|
|
|
|
|
|
return out
|
|
|
|
|