From b082a5b4bf0e65be76584f4aad9b15cb1742740e Mon Sep 17 00:00:00 2001 From: Vladimir Protsenko Date: Mon, 22 Apr 2024 07:05:27 +0000 Subject: [PATCH] update --- src/models/__init__.py | 2 + src/models/munet.py | 65 ----------------- src/models/rcnet.py | 2 +- src/models/{mulut.py => sdylut.py} | 0 src/models/sdynet.py | 108 +++++++++++++++++++++++++++++ 5 files changed, 111 insertions(+), 66 deletions(-) delete mode 100644 src/models/munet.py rename src/models/{mulut.py => sdylut.py} (100%) create mode 100644 src/models/sdynet.py diff --git a/src/models/__init__.py b/src/models/__init__.py index 597a06e..2a089d1 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -2,6 +2,7 @@ from . import rcnet from . import rclut from . import srnet from . import srlut +from . import sdynet import torch import numpy as np from pathlib import Path @@ -16,6 +17,7 @@ AVAILABLE_MODELS = { 'RCNetx1': rcnet.RCNetx1, 'RCLutx1': rclut.RCLutx1, 'RCNetx2': rcnet.RCNetx2, 'RCLutx2': rclut.RCLutx2, 'RCNetx2Centered': rcnet.RCNetx2Centered, 'RCLutx2Centered': rclut.RCLutx2Centered, + 'SDYNetx1': sdynet.SDYNetx1 } def SaveCheckpoint(model, path): diff --git a/src/models/munet.py b/src/models/munet.py deleted file mode 100644 index 1d0cd32..0000000 --- a/src/models/munet.py +++ /dev/null @@ -1,65 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from common.utils import round_func -from common import lut -from pathlib import Path -# from .mulut import MuLutx1, MuLutx2 - -# Huang G. et al. Densely connected convolutional networks //Proceedings of the IEEE conference on computer vision and pattern recognition. – 2017. – С. 4700-4708. -# https://ar5iv.labs.arxiv.org/html/1608.06993 -# https://github.com/andreasveit/densenet-pytorch/blob/63152f4a40644b62717749536ed2e011c6e4d9ab/densenet.py#L40 -class DenseConvUpscaleBlock(nn.Module): - def __init__(self, hidden_dim = 32, layers_count=5, upscale_factor=1): - super(DenseConvUpscaleBlock, self).__init__() - assert layers_count > 0 - self.upscale_factor = upscale_factor - - self.percieve = nn.Conv2d(1, hidden_dim, kernel_size=(2, 2), padding='valid', stride=1, dilation=1, bias=True) - self.convs = [] - for i in range(layers_count): - self.convs.append(nn.Conv2d(in_channels = (i+1)*hidden_dim, out_channels = hidden_dim, kernel_size = 1, stride=1, padding=0, dilation=1, bias=True)) - self.convs = nn.ModuleList(self.convs) - - for name, p in self.named_parameters(): - if "weight" in name: nn.init.kaiming_normal_(p) - if "bias" in name: nn.init.constant_(p, 0) - - self.project_channels = nn.Conv2d(in_channels = (layers_count+1)*hidden_dim, out_channels = upscale_factor * upscale_factor, kernel_size = 1, stride=1, padding=0, dilation=1, bias=True) - self.shuffle = nn.PixelShuffle(upscale_factor) - - def forward(self, x): - x = (x-127.5)/127.5 - x = torch.relu(self.percieve(x)) - for conv in self.convs: - x = torch.cat([x, torch.relu(conv(x))], dim=1) - x = self.shuffle(self.project_channels(x)) - x = torch.tanh(x) - x = round_func(x*127.5 + 127.5) - return x - -class MuNetx1(nn.Module): - def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): - super(SRNet2x2, self).__init__() - self.scale = scale - self.stage = DenseConvUpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) - - def forward(self, x): - b,c,h,w = x.shape - x = x.view(b*c, 1, h, w) - output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) - for rotations_count in range(4): - rotated = torch.rot90(x, k=rotations_count, dims=[2, 3]) - rotated_padded = F.pad(rotated, pad=[0,1,0,1], mode='replicate') - rotated_prediction = self.stage(rotated_padded) - unrotated_prediction = torch.rot90(rotated_prediction, k=-rotations_count, dims=[2, 3]) - output += unrotated_prediction - output /= 4 - output = output.view(b, c, h*self.scale, w*self.scale) - return output - - def get_lut_model(self, quantization_interval=16, batch_size=2**10): - stage_lut = lut.transfer_2x2_input_SxS_output(self.stage, quantization_interval=quantization_interval, batch_size=batch_size) - lut_model = MuLutx1.init_from_lut(stage_lut) - return lut_model \ No newline at end of file diff --git a/src/models/rcnet.py b/src/models/rcnet.py index 3c8dd32..8699b70 100644 --- a/src/models/rcnet.py +++ b/src/models/rcnet.py @@ -48,7 +48,7 @@ class ReconstructedConvCentered(nn.Module): def pixel_wise_forward(self, x): x = (x-127.5)/127.5 - out = torch.einsum('bik,ij,ij -> bik', x, self.projection1, self.projection2) + out = torch.einsum('bwk,wh,wh -> bwk', x, self.projection1, self.projection2) out = torch.tanh(out) out = out*127.5 + 127.5 return out diff --git a/src/models/mulut.py b/src/models/sdylut.py similarity index 100% rename from src/models/mulut.py rename to src/models/sdylut.py diff --git a/src/models/sdynet.py b/src/models/sdynet.py new file mode 100644 index 0000000..9127d99 --- /dev/null +++ b/src/models/sdynet.py @@ -0,0 +1,108 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from common.utils import round_func +from common import lut +from pathlib import Path +from . import sdylut + +# Huang G. et al. Densely connected convolutional networks //Proceedings of the IEEE conference on computer vision and pattern recognition. – 2017. – С. 4700-4708. +# https://ar5iv.labs.arxiv.org/html/1608.06993 +# https://github.com/andreasveit/densenet-pytorch/blob/63152f4a40644b62717749536ed2e011c6e4d9ab/densenet.py#L40 +class DenseConvUpscaleBlock(nn.Module): + def __init__(self, hidden_dim = 32, layers_count=5, upscale_factor=1): + super(DenseConvUpscaleBlock, self).__init__() + assert layers_count > 0 + self.upscale_factor = upscale_factor + self.hidden_dim = hidden_dim + self.percieve = nn.Conv2d(1, hidden_dim, kernel_size=(2, 2), padding='valid', stride=1, dilation=1, bias=True) + self.convs = [] + for i in range(layers_count): + self.convs.append(nn.Conv2d(in_channels = (i+1)*hidden_dim, out_channels = hidden_dim, kernel_size = 1, stride=1, padding=0, dilation=1, bias=True)) + self.convs = nn.ModuleList(self.convs) + + for name, p in self.named_parameters(): + if "weight" in name: nn.init.kaiming_normal_(p) + if "bias" in name: nn.init.constant_(p, 0) + + self.project_channels = nn.Conv2d(in_channels = (layers_count+1)*hidden_dim, out_channels = upscale_factor * upscale_factor, kernel_size = 1, stride=1, padding=0, dilation=1, bias=True) + self.shuffle = nn.PixelShuffle(upscale_factor) + + def forward(self, x): + x = (x-127.5)/127.5 + x = torch.relu(self.percieve(x)) + for conv in self.convs: + x = torch.cat([x, torch.relu(conv(x))], dim=1) + x = self.shuffle(self.project_channels(x)) + x = torch.tanh(x) + x = round_func(x*127.5 + 127.5) + return x + +class PercievePattern(): + def __init__(self, receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]]): + self.receptive_field_idxes = np.array(receptive_field_idxes) + self.window_size = np.max(self.receptive_field_idxes) + 1 + self.receptive_field_idxes = [ + self.receptive_field_idxes[0,0]*self.window_size + self.receptive_field_idxes[0,1], + self.receptive_field_idxes[1,0]*self.window_size + self.receptive_field_idxes[1,1], + self.receptive_field_idxes[2,0]*self.window_size + self.receptive_field_idxes[2,1], + self.receptive_field_idxes[3,0]*self.window_size + self.receptive_field_idxes[3,1], + ] + + def __call__(self, x): + b,c,h,w = x.shape + x = F.pad(x, pad=[0,self.window_size-1,0,self.window_size-1], mode='replicate') + x = F.unfold(input=x, kernel_size=self.window_size) + x = torch.stack([ + x[:,self.receptive_field_idxes[0],:], + x[:,self.receptive_field_idxes[1],:], + x[:,self.receptive_field_idxes[2],:], + x[:,self.receptive_field_idxes[3],:] + ], 2) + x = x.reshape(x.shape[0]*x.shape[1], 1, 2, 2) + return x + + +class SDYNetx1(nn.Module): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SDYNetx1, self).__init__() + self.scale = scale + self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]]) + self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]]) + self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]]) + self.stageS = DenseConvUpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) + self.stageD = DenseConvUpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) + self.stageY = DenseConvUpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) + + def forward(self, x): + b,c,h,w = x.shape + x = x.view(b*c, 1, h, w) + output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) + for rotations_count in range(4): + rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1]) + rb,rc,rh,rw = rotated.shape + + s = self.stageS(self._extract_pattern_S(rotated)) + s = s.view(rb*rc, 1, rh, rw, self.scale, self.scale).permute(0,1,2,4,3,5).reshape(rb*rc, 1, rh*self.scale, rw*self.scale) + s = torch.rot90(s, k=-rotations_count, dims=[-2, -1]) + output += s + + d = self.stageD(self._extract_pattern_D(rotated)) + d = d.view(rb*rc, 1, rh, rw, self.scale, self.scale).permute(0,1,2,4,3,5).reshape(rb*rc, 1, rh*self.scale, rw*self.scale) + d = torch.rot90(d, k=-rotations_count, dims=[-2, -1]) + output += d + + y = self.stageY(self._extract_pattern_Y(rotated)) + y = y.view(rb*rc, 1, rh, rw, self.scale, self.scale).permute(0,1,2,4,3,5).reshape(rb*rc, 1, rh*self.scale, rw*self.scale) + y = torch.rot90(y, k=-rotations_count, dims=[-2, -1]) + output += y + + output /= 4*3 + output = output.view(b, c, h*self.scale, w*self.scale) + return output + + def get_lut_model(self, quantization_interval=16, batch_size=2**10): + stage_lut = lut.transfer_2x2_input_SxS_output(self.stage, quantization_interval=quantization_interval, batch_size=batch_size) + lut_model = sdylut.SDYLutx1.init_from_lut(stage_lut) + return lut_model \ No newline at end of file