63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
from torch import nn as nn
|
|
from torch.nn import functional as F
|
|
|
|
|
|
# https://github.com/xinntao/Real-ESRGAN/blob/master/realesrgan/archs/srvgg_arch.py
|
|
class SRVGGNetCompact(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_in_ch=3,
|
|
num_out_ch=3,
|
|
num_feat=64,
|
|
num_conv=16,
|
|
upscale=4,
|
|
act_type="prelu",
|
|
):
|
|
super(SRVGGNetCompact, self).__init__()
|
|
self.num_in_ch = num_in_ch
|
|
self.num_out_ch = num_out_ch
|
|
self.num_feat = num_feat
|
|
self.num_conv = num_conv
|
|
self.upscale = upscale
|
|
self.act_type = act_type
|
|
|
|
self.body = nn.ModuleList()
|
|
# the first conv
|
|
self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))
|
|
# the first activation
|
|
if act_type == "relu":
|
|
activation = nn.ReLU(inplace=True)
|
|
elif act_type == "prelu":
|
|
activation = nn.PReLU(num_parameters=num_feat)
|
|
elif act_type == "leakyrelu":
|
|
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
|
|
self.body.append(activation)
|
|
|
|
# the body structure
|
|
for _ in range(num_conv):
|
|
self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))
|
|
# activation
|
|
if act_type == "relu":
|
|
activation = nn.ReLU(inplace=True)
|
|
elif act_type == "prelu":
|
|
activation = nn.PReLU(num_parameters=num_feat)
|
|
elif act_type == "leakyrelu":
|
|
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
|
|
self.body.append(activation)
|
|
|
|
# the last conv
|
|
self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))
|
|
# upsample
|
|
self.upsampler = nn.PixelShuffle(upscale)
|
|
|
|
def forward(self, x):
|
|
out = x
|
|
for i in range(0, len(self.body)):
|
|
out = self.body[i](out)
|
|
|
|
out = self.upsampler(out)
|
|
# add the nearest upsampled image, so that the network learns the residual
|
|
base = F.interpolate(x, scale_factor=self.upscale, mode="nearest")
|
|
out += base
|
|
return out
|