diff --git a/src/common/layers.py b/src/common/layers.py index df84314..95b3731 100644 --- a/src/common/layers.py +++ b/src/common/layers.py @@ -25,7 +25,7 @@ class PercievePattern(): return x class UpscaleBlock(nn.Module): - def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1): + def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255): super(UpscaleBlock, self).__init__() assert layers_count > 0 self.upscale_factor = upscale_factor @@ -38,15 +38,18 @@ class UpscaleBlock(nn.Module): self.linear_projections = nn.ModuleList(self.linear_projections) self.project_channels = nn.Linear(in_features=(layers_count+1)*hidden_dim, out_features=upscale_factor * upscale_factor, bias=True) + + self.in_bias = self.in_scale = input_max_value/2 + self.out_bias = self.out_scale = output_max_value/2 def forward(self, x): - x = (x-127.5)/127.5 + x = (x-self.in_bias)/self.in_scale x = torch.relu(self.embed(x)) for linear_projection in self.linear_projections: x = torch.cat([x, torch.relu(linear_projection(x))], dim=2) x = self.project_channels(x) x = torch.tanh(x) - x = x*127.5 + 127.5 + x = x*self.out_scale + self.out_bias return x class RgbToYcbcr(nn.Module): diff --git a/src/common/losses.py b/src/common/losses.py index 688595b..69f9aab 100644 --- a/src/common/losses.py +++ b/src/common/losses.py @@ -1,5 +1,6 @@ import torch +from torch import nn class FourierLoss(nn.Module): def __init__(self, weight=None, size_average=True): diff --git a/src/models/srnet.py b/src/models/srnet.py index 798cbda..eb818be 100644 --- a/src/models/srnet.py +++ b/src/models/srnet.py @@ -7,7 +7,6 @@ from common import lut from pathlib import Path from . import srlut from common import layers -from itertools import cycle from common import losses class SRNet(nn.Module): @@ -189,6 +188,79 @@ class SRNetR90Y(nn.Module): return F.mse_loss(pred/255, target/255) return loss_fn + +class SRMsbLsbR90Net(nn.Module): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SRMsbLsbR90Net, self).__init__() + self.scale = scale + self.hidden_dim = hidden_dim + self.layers_count = layers_count + + self.msb_fn = layers.UpscaleBlock( + in_features=4, + hidden_dim=hidden_dim, + layers_count=layers_count, + upscale_factor=self.scale, + input_max_value=255, + output_max_value=15 + ) + self.lsb_fn = layers.UpscaleBlock( + in_features=4, + hidden_dim=hidden_dim, + layers_count=layers_count, + upscale_factor=self.scale, + input_max_value=15, + output_max_value=15 + ) + self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) + + def forward_stage(self, x, scale, percieve_pattern, stage): + b,c,h,w = x.shape + x = percieve_pattern(x) + x = stage(x) + x = round_func(x) + x = x.reshape(b, c, h, w, scale, scale) + x = x.permute(0,1,2,4,3,5) + x = x.reshape(b, c, h*scale, w*scale) + return x + + def forward(self, x, config=None): + b,c,h,w = x.shape + x = x.reshape(b*c, 1, h, w) + + lsb = x % 16 + msb = x - lsb + + output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) + output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) + for rotations_count in range(4): + rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3]) + rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3]) + output_msb_r = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_S, msb_fn) + output_lsb_r = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_S, lsb_fn) + output_msb_r = round_func(output_msb_r) * 15 + output_lsb_r = round_func(output_lsb_r) + output_msb += torch.rot90(output_msb_r, k=-rotations_count, dims=[2, 3]) + output_lsb += torch.rot90(output_lsb_r, k=-rotations_count, dims=[2, 3]) + output_msb /= 4 + output_lsb /= 4 + if not config is None and config.current_iter % config.display_step == 0: + config.writer.add_histogram('output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter) + config.writer.add_histogram('output_msb', output_msb.detach().cpu().numpy(), config.current_iter) + x = output_msb + output_lsb + x = x.reshape(b, c, h*self.scale, w*self.scale) + return x + + def get_lut_model(self, quantization_interval=16, batch_size=2**10): + raise NotImplementedError + + def get_loss_fn(self): + fourier_loss_fn = losses.FocalFrequencyLoss() + def loss_fn(pred, target): + return fourier_loss_fn(pred, target) + return loss_fn + + class SRMsbLsb4R90Net(nn.Module): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SRMsbLsb4R90Net, self).__init__() @@ -200,13 +272,17 @@ class SRMsbLsb4R90Net(nn.Module): in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, - upscale_factor=self.scale + upscale_factor=self.scale, + input_max_value=255, + output_max_value=15 ) for x in range(4)]) self.lsb_fns = nn.ModuleList([layers.UpscaleBlock( in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, - upscale_factor=self.scale + upscale_factor=self.scale, + input_max_value=15, + output_max_value=15 ) for x in range(4)]) self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) @@ -229,13 +305,13 @@ class SRMsbLsb4R90Net(nn.Module): output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) - for rotations_count, msb_fn, lsb_fn in zip(range(4), cycle(self.msb_fns), cycle(self.lsb_fns)): + for rotations_count, msb_fn, lsb_fn in zip(range(4), self.msb_fns, self.lsb_fns): rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3]) rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3]) output_msb_r = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_S, msb_fn) output_lsb_r = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_S, lsb_fn) - output_msb_r = round_func((output_msb_r / 255)*16) * 15 - output_lsb_r = (output_lsb_r / 255) * 15 + output_msb_r = round_func(output_msb_r) * 15 + output_lsb_r = round_func(output_lsb_r) output_msb += torch.rot90(output_msb_r, k=-rotations_count, dims=[2, 3]) output_lsb += torch.rot90(output_lsb_r, k=-rotations_count, dims=[2, 3]) output_msb /= 4 @@ -253,5 +329,5 @@ class SRMsbLsb4R90Net(nn.Module): def get_loss_fn(self): fourier_loss_fn = losses.FocalFrequencyLoss() def loss_fn(pred, target): - return fourier_loss_fn(pred/255, target/255) * 1e8 + return fourier_loss_fn(pred, target) return loss_fn \ No newline at end of file diff --git a/src/train.py b/src/train.py index 5f12e08..2eb6fee 100644 --- a/src/train.py +++ b/src/train.py @@ -113,7 +113,6 @@ if __name__ == "__main__": model = AVAILABLE_MODELS[config.model]( quantization_interval = 2**(8-config.quantization_bits), scale = config.scale) model = model.to(torch.device(config.device)) optimizer = AdamWScheduleFree(model.parameters(), lr=1e-2, betas=(0.9, 0.95)) - # optimizer = optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95)) print(optimizer) prepare_experiment_folder(config)