def __init__(self, num_blocks=None, num_channel=64, block=ResBlock):
super().__init__()
if num_blocks is None:
num_blocks = [5, 10, 15]
assert len(num_blocks) == 3
self.init = nn.Sequential(
MeanShift(sign=-1),
ConvolutionBlock(in_channels=3, out_channels=num_channel))
self.model_1 = nn.Sequential(
*tuple([
block(in_channels=num_channel) for _ in range(num_blocks[0])
]), ConvolutionDownscale(channels=num_channel),
ConvolutionDownscale(channels=num_channel))
self.model_2 = nn.Sequential(
*tuple([
block(in_channels=num_channel) for _ in range(num_blocks[1])
]), ConvolutionDownscale(channels=num_channel))
self.model_3 = nn.Sequential(*tuple(
[block(in_channels=num_channel) for _ in range(num_blocks[2])]))
self.downscale_1 = ConvolutionDownscale(channels=num_channel)
self.downscale_2 = ConvolutionDownscale(channels=num_channel)
self.end = nn.Sequential(
ConvolutionBlock(in_channels=num_channel * 3, out_channels=3),
MeanShift(sign=1))
def __init__(self, num_blocks=32, num_channels=256):
super().__init__()
self.conv_in = ConvolutionBlock(in_channels=3, out_channels=num_channels, kernel_size=1, padding=0)
res_blocks = [Res2Block(in_channels=num_channels, up_channels=num_channels, out_channels=num_channels) for i in range(num_blocks)]
self.res_blocks = nn.Sequential(*res_blocks)
self.upsample = TransposeUpscale(channels=num_channels*2, scale=4, mode=2)
self.conv_out = ConvolutionBlock(in_channels=num_channels*2, out_channels=3)
def __init__(self, num_blocks=50, channel=128):
super().__init__()
self.model = nn.Sequential(
ConvolutionBlock(in_channels=3, out_channels=channel),
*tuple([ResBlock(in_channels=channel) for _ in range(num_blocks)]),
ConvolutionBlock(in_channels=channel)
)
def __init__(self, block_sequence=None, num_channel=256,
rgb_mean=(0.4488, 0.4371, 0.4040),
rgb_std=(1.0, 1.0, 1.0), rgb_range=255):
super().__init__()
if block_sequence is None:
block_sequence = [15, 10, 5]
self.init = nn.Sequential(
MeanShift(rgb_range, rgb_mean, rgb_std, sign=-1),
ConvolutionBlock(in_channels=3, out_channels=num_channel)
)
self.skip_upscale = nn.Sequential(
TransposeUpscale(channels=num_channel),
ConvolutionBlock(in_channels=num_channel, out_channels=3)
)
model_body = []
for seq in block_sequence:
model_body.extend([ResBlock(in_channels=num_channel) for _ in range(seq)])
model_body.append(PixelShuffleUpscale(channels=num_channel))
model_body.append(ConvolutionBlock(in_channels=num_channel, out_channels=3))
self.blocks = nn.Sequential(*tuple(model_body))
self.model_end = nn.Sequential(
ConvolutionBlock(in_channels=num_channel, out_channels=3),
MeanShift(rgb_range, rgb_mean, rgb_std, sign=1)
)
def __init__(self, in_channels):
super().__init__()
self.preproc = nn.Sequential(
ConvolutionBlock(in_channels=in_channels,
out_channels=4,
activation=nn.ReLU))
self.postproc = nn.Sequential(
ConvolutionBlock(in_channels=4, out_channels=1, activation=None),
nn.Sigmoid())
def __init__(self, num_res=30, channel=128):
super().__init__()
self.model = nn.Sequential(
ConvolutionBlock(in_channels=3, out_channels=channel),
*tuple([ResBlock(in_channels=channel) for _ in range(num_res)]),
ConvolutionBlock(in_channels=channel),
PixelShuffleUpscale(channels=channel),
PixelShuffleUpscale(channels=channel),
ConvolutionBlock(in_channels=channel, out_channels=3)
)
def __init__(self, num_se=30, channel=128, bias=True):
super().__init__()
self.model = nn.Sequential(
ConvolutionBlock(in_channels=3, out_channels=channel),
*tuple([
ResBlock(in_channels=channel),
ChannelAttentionBlock(channel=channel)
] for _ in range(num_se)),
ConvolutionBlock(in_channels=channel, out_channels=channel),
PixelShuffleUpscale(channels=channel),
PixelShuffleUpscale(channels=channel),
ConvolutionBlock(in_channels=channel, out_channels=3))
def __init__(self, scale, num_blocks=32, num_channel=256):
super().__init__()
self.model_0 = nn.Sequential(
MeanShift(sign=-1),
ConvolutionBlock(in_channels=3, out_channels=num_channel)
)
self.model_1 = nn.Sequential(*tuple([ResBlock(num_channel) for _ in range(num_blocks)]))
self.model_2 = nn.Sequential(
*tuple([PixelShuffleUpscale(channels=num_channel) for _ in range(int(log2(scale)))]),
ConvolutionBlock(in_channels=num_channel, out_channels=3),
MeanShift(sign=1)
)
def __init__(self, scale=4, num_groups=12, num_channel=128, rep_pad=False):
super().__init__()
self.model_0 = nn.Sequential(
MeanShift(sign=-1),
ConvolutionBlock(in_channels=3, out_channels=num_channel, rep_pad=rep_pad)
)
self.model_1 = nn.Sequential(*tuple([BasicGroup(num_channel, rep_pad=rep_pad) for _ in range(num_groups)]))
upscaler = [PixelShuffleUpscale(channels=num_channel,
basic_block=DepthSeparableConvBlock, rep_pad=rep_pad) for _ in range(int(log2(scale)))]
self.model_2 = nn.Sequential(
*tuple(upscaler),
ConvolutionBlock(in_channels=num_channel, out_channels=3, rep_pad=rep_pad),
MeanShift(sign=1)
)
def __init__(self, num_channel, num_block=5, res_scale=1, rep_pad=False):
super().__init__()
self.model_body = nn.Sequential(
*tuple([BasicBlock(num_channel, rep_pad=rep_pad) for _ in range(num_block)]),
ConvolutionBlock(in_channels=num_channel, rep_pad=rep_pad)
)
self.res_scale = res_scale
def __init__(self, scale=4, num_groups=20, num_channel=64):
super().__init__()
self.model_0 = nn.Sequential(
MeanShift(sign=-1),
ConvolutionBlock(in_channels=3, out_channels=num_channel))
self.model_1 = nn.Sequential(
*tuple([BasicGroup(num_channel) for _ in range(num_groups)]))
upscaler = [
PixelShuffleUpscale(channels=num_channel)
for _ in range(int(log2(scale)))
]
self.model_2 = nn.Sequential(
*tuple(upscaler),
ConvolutionBlock(in_channels=num_channel, out_channels=3),
MeanShift(sign=1))
def __init__(self,
channels,
scale=2,
activation=nn.LeakyReLU,
batch_norm=None,
mode=1,
rep_pad=False):
super().__init__()
if mode == 1:
k_size, p_size = 4, 1
elif mode == 2:
k_size, p_size = 6, 2
else:
raise ValueError('wrong mode parameters')
model_body = []
for _ in range(int(log2(scale))):
model_body.append(ConvolutionBlock(channels))
model_body.append(
nn.ConvTranspose2d(in_channels=channels,
out_channels=channels,
kernel_size=k_size,
padding=p_size,
stride=2,
rep_pad=rep_pad))
if activation is not None:
model_body.append(activation(True))
if batch_norm is not None:
model_body.append(batch_norm(channels))
self.model = nn.Sequential(*tuple(model_body))
def __init__(self, num_groups=None, num_channel=128):
super().__init__()
if num_groups is None:
num_groups = [6, 3]
self.model_0 = nn.Sequential(
MeanShift(sign=-1),
ConvolutionBlock(in_channels=3, out_channels=num_channel)
)
self.model_1 = nn.Sequential(*tuple([BasicGroup(num_channel) for _ in range(num_groups[0])]))
self.upscale_1 = PixelShuffleUpscale(channels=num_channel)
self.model_2 = nn.Sequential(*tuple([BasicGroup(num_channel) for _ in range(num_groups[1])]))
self.upscale_2 = PixelShuffleUpscale(channels=num_channel)
self.model_3 = nn.Sequential(
ConvolutionBlock(in_channels=num_channel, out_channels=3),
MeanShift(sign=1)
)
def __init__(self, channels, out_channels=None, scale=2):
super().__init__()
out_channels = channels if out_channels is None else out_channels
filter_size = scale * 4
stride = scale
pad = scale * 3
pad_top, pad_left = pad // 2, pad // 2
pad_bottom, pad_right = pad - pad_top, pad - pad_left
self.model = nn.Sequential(
nn.ReplicationPad2d((pad_left, pad_right, pad_top, pad_bottom)),
nn.Conv2d(in_channels=channels,
out_channels=out_channels,
kernel_size=filter_size,
stride=stride), nn.ReLU(inplace=True),
ConvolutionBlock(in_channels=channels))
def __init__(self):
super().__init__()
self.model = nn.Sequential(
ChannelPool(),
ConvolutionBlock(in_channels=3, out_channels=1, activation=None),
nn.Sigmoid())
