diff --git a/src/common/layers.py b/src/common/layers.py index 47b8b6e..7c75f65 100644 --- a/src/common/layers.py +++ b/src/common/layers.py @@ -361,7 +361,7 @@ class ChebyKANLayer(nn.Module): """ https://github.com/SynodicMonth/ChebyKAN/blob/main/ChebyKANLayer.py """ - def __init__(self, in_features, out_features, degree=5): + def __init__(self, in_features, out_features, degree=8): super(ChebyKANLayer, self).__init__() self.inputdim = in_features self.outdim = out_features @@ -396,7 +396,7 @@ class ChebyKANLayer(nn.Module): class UpscaleBlockChebyKAN(nn.Module): - def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255): + def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255, degree=8): super(UpscaleBlockChebyKAN, self).__init__() assert layers_count > 0 self.upscale_factor = upscale_factor @@ -405,7 +405,7 @@ class UpscaleBlockChebyKAN(nn.Module): self.linear_projections = [] for i in range(layers_count): - self.linear_projections.append(ChebyKANLayer(in_features=hidden_dim, out_features=hidden_dim)) + self.linear_projections.append(ChebyKANLayer(in_features=hidden_dim, out_features=hidden_dim, degree=degree)) self.linear_projections = nn.ModuleList(self.linear_projections) self.project_channels = nn.Linear(in_features=hidden_dim, out_features=upscale_factor * upscale_factor, bias=True) diff --git a/src/models/hdbnet.py b/src/models/hdbnet.py index 4f5bacf..15fd3a7 100644 --- a/src/models/hdbnet.py +++ b/src/models/hdbnet.py @@ -11,10 +11,11 @@ from itertools import cycle from models.base import SRNetBase class HDBNet(SRNetBase): - def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4, rotations = 4): super(HDBNet, self).__init__() assert scale == 4 self.scale = scale + self.rotations = rotations self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) @@ -31,9 +32,7 @@ class HDBNet(SRNetBase): self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3) self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3) self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2) - self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2) - - self.rotations = 4 + self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2) def forward_stage(self, x, scale, percieve_pattern, stage): b,c,h,w = x.shape @@ -113,6 +112,102 @@ class HDBNet(SRNetBase): return F.mse_loss(pred/255, target/255) return loss_fn +class HDBNetv2(SRNetBase): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4, rotations = 4): + super(HDBNetv2, self).__init__() + assert scale == 4 + self.scale = scale + self.rotations = rotations + self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) + self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) + self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) + self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255) + self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255) + + self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) + self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) + self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255) + self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255) + self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255) + + self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3) + self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3) + self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3) + self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2) + self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2) + + 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*2, w*2], dtype=x.dtype, device=x.device) + output_lsb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device) + for rotations_count in range(self.rotations): + rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3]) + rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3]) + output_msbt = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage1_3H) + \ + self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage1_3D) + \ + self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage1_3B) + output_lsbt = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage1_2H) + \ + self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage1_2D) + if not config is None and config.current_iter % config.display_step == 0: + config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter) + config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) + output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3]) + output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3]) + output_msb /= self.rotations*3 + output_lsb /= self.rotations*2 + output = output_msb + output_lsb + x = output.clamp(0, 255) + lsb = x % 16 + msb = x - lsb + output_msb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device) + output_lsb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device) + for rotations_count in range(self.rotations): + rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3]) + rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3]) + output_msbt = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H) + \ + self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D) + \ + self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B) + output_lsbt = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H) + \ + self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D) + if not config is None and config.current_iter % config.display_step == 0: + config.writer.add_histogram('s2_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter) + config.writer.add_histogram('s2_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) + output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3]) + output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3]) + output_msb /= self.rotations*3 + output_lsb /= self.rotations*2 + output = output_msb + output_lsb + x = output.clamp(0, 255) + 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): + stage1_3H = lut.transfer_3_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size) + stage1_3D = lut.transfer_3_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size) + stage1_3B = lut.transfer_3_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size) + stage1_2H = lut.transfer_2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size) + stage1_2D = lut.transfer_2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size) + + stage2_3H = lut.transfer_3_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size) + stage2_3D = lut.transfer_3_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size) + stage2_3B = lut.transfer_3_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size) + stage2_2H = lut.transfer_2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size) + stage2_2D = lut.transfer_2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size) + + lut_model = hdblut.HDBLut.init_from_numpy( + stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D, + stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D + ) + return lut_model + + def get_loss_fn(self): + def loss_fn(pred, target): + return F.mse_loss(pred/255, target/255) + return loss_fn + class HDBLNet(nn.Module): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): diff --git a/src/models/sdynet.py b/src/models/sdynet.py index b37e5df..745ce07 100644 --- a/src/models/sdynet.py +++ b/src/models/sdynet.py @@ -7,9 +7,9 @@ from common import lut from common import layers from pathlib import Path from . import sdylut +from models.base import SRNetBase - -class SDYNetx1(nn.Module): +class SDYNetx1(SRNetBase): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SDYNetx1, self).__init__() self.scale = scale @@ -20,16 +20,6 @@ class SDYNetx1(nn.Module): self.stage1_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) self.stage1_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) - 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) @@ -39,7 +29,6 @@ class SDYNetx1(nn.Module): output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage1_Y) output /= 3 x = output - x = round_func(x) x = x.reshape(b, c, h*self.scale, w*self.scale) return x @@ -56,7 +45,7 @@ class SDYNetx1(nn.Module): return loss_fn -class SDYNetx2(nn.Module): +class SDYNetx2(SRNetBase): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SDYNetx2, self).__init__() self.scale = scale @@ -70,16 +59,6 @@ class SDYNetx2(nn.Module): self.stage2_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) self.stage2_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) - 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) @@ -89,14 +68,12 @@ class SDYNetx2(nn.Module): output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y) output /= 3 x = output - x = round_func(x) output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S) output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D) output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage2_Y) output /= 3 x = output - x = round_func(x) x = x.reshape(b, c, h*self.scale, w*self.scale) return x @@ -115,7 +92,7 @@ class SDYNetx2(nn.Module): return F.mse_loss(pred/255, target/255) return loss_fn -class SDYNetx3(nn.Module): +class SDYNetx3(SRNetBase): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SDYNetx3, self).__init__() self.scale = scale @@ -132,16 +109,6 @@ class SDYNetx3(nn.Module): self.stage3_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) self.stage3_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) - 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) @@ -151,21 +118,18 @@ class SDYNetx3(nn.Module): output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y) output /= 3 x = output - x = round_func(x) output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device) output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage2_S) output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage2_D) output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage2_Y) output /= 3 x = output - x = round_func(x) output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage3_S) output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage3_D) output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage3_Y) output /= 3 x = output - x = round_func(x) x = x.reshape(b, c, h*self.scale, w*self.scale) return x @@ -222,7 +186,6 @@ class SDYNetR90x1(nn.Module): output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1]) output /= 4*3 x = output - x = round_func(x) x = x.reshape(b, c, h*self.scale, w*self.scale) return x @@ -238,7 +201,7 @@ class SDYNetR90x1(nn.Module): return F.mse_loss(pred/255, target/255) return loss_fn -class SDYNetR90x2(nn.Module): +class SDYNetR90x2(SRNetBase): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SDYNetR90x2, self).__init__() self.scale = scale @@ -252,16 +215,6 @@ class SDYNetR90x2(nn.Module): self.stage2_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) self.stage2_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) - 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.view(b*c, 1, h, w) @@ -275,7 +228,7 @@ class SDYNetR90x2(nn.Module): output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1]) output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1]) output_1 /= 4*3 - x = round_func(output_1) + x = output_1 output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) output_2 += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S) output_2 += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D) @@ -286,7 +239,7 @@ class SDYNetR90x2(nn.Module): output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_D, self.stage2_D), k=-rotations_count, dims=[-2, -1]) output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage2_Y), k=-rotations_count, dims=[-2, -1]) output_2 /= 4*3 - x = round_func(output_2) + x = output_2 x = x.view(b, c, h*self.scale, w*self.scale) return x