updates

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):

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):

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

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)