models/team23_DSCF.py

from collections import OrderedDict
import torch
from torch import nn as nn
import torch.nn.functional as F
import math
from typing import Optional, List
# from IPython import embed

class LoRALayer():
    def __init__(
        self, 
        r: int, 
        lora_alpha: int, 
        lora_dropout: float,
        merge_weights: bool,
    ):
        self.r = r
        self.lora_alpha = lora_alpha
        # Optional dropout
        if lora_dropout > 0.:
            self.lora_dropout = nn.Dropout(p=lora_dropout)
        else:
            self.lora_dropout = lambda x: x
        # Mark the weight as unmerged
        self.merged = False
        self.merge_weights = merge_weights

class Lora_Conv2d(nn.Conv2d, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
        self, 
        in_channels: int, 
        out_channels: int,
        kernel_size: int,
        r: int = 0, 
        lora_alpha: int = 1, 
        lora_dropout: float = 0.,
        merge_weights: bool = True,
        **kwargs
    ):
        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                           merge_weights=merge_weights)
        assert type(kernel_size) is int
        # print("in init")
        # embed()
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(
                self.weight.new_zeros((r*kernel_size, in_channels*kernel_size))
            )
            self.lora_B = nn.Parameter(
                self.weight.new_zeros((out_channels*kernel_size, r*kernel_size))
            )
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        # Freeze the bias
        # if self.bias is not None:
        #     self.bias.requires_grad = False
        self.reset_parameters()

    def reset_parameters(self):
        nn.Conv2d.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize A the same way as the default for nn.Linear and B to zero
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def train(self, mode: bool = True): # True for train and False for eval
 
        nn.Conv2d.train(self, mode)
        if mode:
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                self.weight.data -= (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling
                self.merged = False
        else:
            # print("test")
            # embed()
            if self.merge_weights and not self.merged:
                # print("merging")
                # embed()
                # Merge the weights and mark it
                self.weight.data += (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling
                self.merged = True

    def forward(self, x: torch.Tensor):
        # print(f"LoRA merged status: {self.merged}")
        if self.r > 0 and not self.merged:
            # print(f"lora_A: {self.lora_A}")
            # print(f"lora_B: {self.lora_B}")
            # print(f"LoRA contribution: {(self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling}")
            return F.conv2d(
                x, 
                self.weight + (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling,
                self.bias, self.stride, self.padding, self.dilation, self.groups
            )
        
        return nn.Conv2d.forward(self, x)
def _make_pair(value):
    if isinstance(value, int):
        value = (value,) * 2
    return value


def conv_layer(in_channels,
               out_channels,
               kernel_size,
               bias=True):
    """
    Re-write convolution layer for adaptive `padding`.
    """
    kernel_size = _make_pair(kernel_size)
    padding = (int((kernel_size[0] - 1) / 2),
               int((kernel_size[1] - 1) / 2))
    return nn.Conv2d(in_channels,
                     out_channels,
                     kernel_size,
                     padding=padding,
                     bias=bias)


def activation(act_type, inplace=True, neg_slope=0.05, n_prelu=1):
    """
    Activation functions for ['relu', 'lrelu', 'prelu'].
    Parameters
    ----------
    act_type: str
        one of ['relu', 'lrelu', 'prelu'].
    inplace: bool
        whether to use inplace operator.
    neg_slope: float
        slope of negative region for `lrelu` or `prelu`.
    n_prelu: int
        `num_parameters` for `prelu`.
    ----------
    """
    act_type = act_type.lower()
    if act_type == 'relu':
        layer = nn.ReLU(inplace)
    elif act_type == 'lrelu':
        layer = nn.LeakyReLU(neg_slope, inplace)
    elif act_type == 'prelu':
        layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope)
    else:
        raise NotImplementedError(
            'activation layer [{:s}] is not found'.format(act_type))
    return layer


def sequential(*args):
    """
    Modules will be added to the a Sequential Container in the order they
    are passed.

    Parameters
    ----------
    args: Definition of Modules in order.
    -------
    """
    if len(args) == 1:
        if isinstance(args[0], OrderedDict):
            raise NotImplementedError(
                'sequential does not support OrderedDict input.')
        return args[0]
    modules = []
    for module in args:
        if isinstance(module, nn.Sequential):
            for submodule in module.children():
                modules.append(submodule)
        elif isinstance(module, nn.Module):
            modules.append(module)
    return nn.Sequential(*modules)


def pixelshuffle_block(in_channels,
                       out_channels,
                       upscale_factor=2,
                       kernel_size=3):
    """
    Upsample features according to `upscale_factor`.
    """
    conv = conv_layer(in_channels,
                      out_channels * (upscale_factor ** 2),
                      kernel_size)
    pixel_shuffle = nn.PixelShuffle(upscale_factor)
    return sequential(conv, pixel_shuffle)

class Conv3XC(nn.Module):
    def __init__(self, c_in, c_out, gain1=1, gain2=0, s=1, bias=True, relu=False):
        super(Conv3XC, self).__init__()
        self.weight_concat = None
        self.bias_concat = None
        self.update_params_flag = False
        self.stride = s
        self.has_relu = relu
  

        self.eval_conv = nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=3, padding=1, stride=s, bias=bias)

    def forward(self, x):
        out = self.eval_conv(x)
        if self.has_relu:
            out = F.leaky_relu(out, negative_slope=0.05)
        return out 


class SPAB(nn.Module):
    def __init__(self,
                 in_channels,
                 mid_channels=None,
                 out_channels=None,
                 bias=False):
        super(SPAB, self).__init__()
        if mid_channels is None:
            mid_channels = in_channels
        if out_channels is None:
            out_channels = in_channels

        self.in_channels = in_channels
        self.c1_r = Conv3XC(in_channels, mid_channels, gain1=2, s=1)
        self.c2_r = Conv3XC(mid_channels, mid_channels, gain1=2, s=1)
        self.c3_r = Conv3XC(mid_channels, out_channels, gain1=2, s=1)
        self.act1 = torch.nn.SiLU(inplace=True)
        # self.act2 = activation('lrelu', neg_slope=0.1, inplace=True)

    def forward(self, x):
        out1 = (self.c1_r(x))
        out1_act = self.act1(out1)

        out2 = (self.c2_r(out1_act))
        out2_act = self.act1(out2)

        out3 = (self.c3_r(out2_act))

        sim_att = torch.sigmoid(out3) - 0.5
        out = (out3 + x) * sim_att
        # out = out3 * sim_att
        # return out, out1, sim_att
        return out, out1, out2,out3


class DSCF(nn.Module):
    """
    Swift Parameter-free Attention Network for Efficient Super-Resolution
    """

    def __init__(self,
                 num_in_ch,
                 num_out_ch,
                 feature_channels=26,
                 upscale=4,
                 bias=True,
                 img_range=255.,
                 rgb_mean=(0.4488, 0.4371, 0.4040)
                 ):
        super(DSCF, self).__init__()

        in_channels = num_in_ch
        out_channels = num_out_ch
        self.img_range = img_range
        self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)

        self.conv_1 = Conv3XC(in_channels, feature_channels, gain1=2, s=1)
        self.block_1 = SPAB(feature_channels, bias=bias)
        self.block_2 = SPAB(feature_channels, bias=bias)
        self.block_3 = SPAB(feature_channels, bias=bias)
        self.block_4 = SPAB(feature_channels, bias=bias)
        self.block_5 = SPAB(feature_channels, bias=bias)
        self.block_6 = SPAB(feature_channels, bias=bias)

        self.conv_cat = conv_layer(feature_channels * 4, feature_channels, kernel_size=1, bias=True)
        self.conv_2 = Conv3XC(feature_channels, feature_channels, gain1=2, s=1)
        
        self.upsampler = pixelshuffle_block(feature_channels, out_channels, upscale_factor=upscale)
        
        # 指定需要替换 LoRA 层的子模块名称
        # desired_submodules = ["conv_1.eval_conv",
        #                       "block_1.c1_r.eval_conv","block_1.c2_r.eval_conv","block_1.c3_r.eval_conv",
        #                       "block_2.c1_r.eval_conv","block_2.c2_r.eval_conv","block_2.c3_r.eval_conv",
        #                       "block_3.c1_r.eval_conv","block_3.c2_r.eval_conv","block_3.c3_r.eval_conv",
        #                       "block_4.c1_r.eval_conv","block_4.c2_r.eval_conv","block_4.c3_r.eval_conv",
        #                       "block_5.c1_r.eval_conv","block_5.c2_r.eval_conv","block_5.c3_r.eval_conv",
        #                       "block_6.c1_r.eval_conv","block_6.c2_r.eval_conv","block_6.c3_r.eval_conv",
        #                       "conv_2.eval_conv",
        #                       "conv_cat", 
        #                       "upsampler.0"]
        
        # desired_submodules = ["conv_2.eval_conv","upsampler.0"]
        # # 替换需要 LoRA 处理的层
        # self.replace_layers(desired_submodules)
        
        # self.mark_only_lora_as_trainable(bias='none')
        # 分层LoRA配置字典（模块名: (r, lora_alpha)）
        # self.lora_config = {
        #     # 高频重建核心层 (最高优先级)
        #     "conv_2.eval_conv": (8, 16),  # 最大秩
        #     "upsampler.0": (8, 16),       # 高秩
            
        #     # 中间处理层 (梯度传播关键路径)
        #     **{f"block_{i}.c{j}_r.eval_conv": (2, 4) 
        #     for i in [2,3,4,5]          # block_2到block_5
        #     for j in [1,2,3]},          # 每个block的三个卷积
            
        #     # 首尾层 (适度调整)
        #     "block_1.c1_r.eval_conv": (2, 4),
        #     "block_1.c2_r.eval_conv": (2, 4),
        #     "block_1.c3_r.eval_conv": (2, 4),
        #     "block_6.c1_r.eval_conv": (2, 4),
        #     "block_6.c2_r.eval_conv": (2, 4),
        #     "block_6.c3_r.eval_conv": (2, 4),
        # }
        
        # # 替换需要 LoRA 处理的层
        # self.replace_layers_with_strategy()
        
        # 冻结非LoRA参数
        # self.mark_only_lora_as_trainable(bias='none')
        # self.cuda()(torch.randn(1, 3, 256, 256).cuda())
        # self.eval().cuda()
        self.eval().cuda()
        input_tensor = torch.randn(1, 3, 256, 256).cuda()
        output = self(input_tensor)
        # 确保 LoRA 层参数可训练
        # print("可训练参数:")
        # for name, param in self.named_parameters():
        #     if param.requires_grad:
        #         print(f"{name}: {param.shape}")
                

    # def replace_layers_with_strategy(self):
    #     """根据分层策略替换卷积层"""
    #     for full_name, (r, alpha) in self.lora_config.items():
    #         parent, child_name = self._get_parent_and_child(full_name)
    #         if parent is None:
    #             # print(f"⚠️ Skip {full_name}: module not found")
    #             continue
                
    #         original_conv = getattr(parent, child_name, None)
    #         if not isinstance(original_conv, nn.Conv2d):
    #             # print(f"⚠️ {full_name} is not Conv2d (found {type(original_conv)})")
    #             continue
                
    #         # 动态设置参数
    #         new_layer = Lora_Conv2d(
    #             in_channels=original_conv.in_channels,
    #             out_channels=original_conv.out_channels,
    #             kernel_size=original_conv.kernel_size[0],
    #             stride=original_conv.stride,
    #             padding=original_conv.padding,
    #             bias=original_conv.bias is not None,
    #             r=r,  # 动态设置秩
    #             lora_alpha=alpha  # 动态设置缩放系数
    #         )
            
    #         # 继承原始权重
    #         with torch.no_grad():
    #             new_layer.weight.copy_(original_conv.weight)
    #             if original_conv.bias is not None:
    #                 new_layer.bias.copy_(original_conv.bias)
            
    #         setattr(parent, child_name, new_layer)
    #         # print(f"✅ {full_name} => r={r}, alpha={alpha}")

    # def _get_parent_and_child(self, module_name):
    #     """
    #     获取模块的父级模块和子模块名称
    #     例如：
    #     module_name = "block_5.c1_r.eval_conv"
    #     则返回 (model.block_5.c1_r, "eval_conv")
    #     """
    #     parts = module_name.split(".")
    #     parent = self
    #     for part in parts[:-1]:  # 遍历到倒数第二个
    #         if hasattr(parent, part):
    #             parent = getattr(parent, part)
    #         else:
    #             return None, None  # 没找到路径
    #     return parent, parts[-1]  # 返回父模块和子模块名称

    # def replace_layers(self, desired_submodules):
    #     """
    #     遍历模型的子模块，将符合条件的层替换为 Lora_Conv2d
    #     """
    #     # 替换conv_layer
    #     for name, module in self._modules.items():
    #         if name in desired_submodules:
    #             print('--------------------self._modules.items--------------------------')
    #             print(name)
    #         if isinstance(module, nn.Conv2d):
    #             print(f"Replacing {name} with Lora_Conv2d")
    #             setattr(self, name, Lora_Conv2d(
    #                 module.in_channels,
    #                 module.out_channels,
    #                 kernel_size=module.kernel_size[0],
    #                 stride=module.stride,
    #                 padding=module.padding,
    #                 bias=True,
    #                 r=2,
    #                 lora_alpha=2
    #             ))
                                    
    # def mark_only_lora_as_trainable(self, bias: str = 'none'):
    #     """
    #     只训练 LoRA 相关参数，而冻结所有其他参数。
        
    #     参数:
    #     - bias: 'none' (不训练 bias), 'all' (训练所有 bias), 'lora_only' (只训练 LoRA 层的 bias)
    #     """
    #     # 冻结所有非 LoRA 参数
    #     # for n, p in self.named_parameters():
    #     #     if 'lora_' not in n:
    #     #         p.requires_grad = False
    #     for n, p in self.named_parameters():
    #         if 'lora_' not in n:  
    #             p.requires_grad = False  # 冻结非 LoRA 参数
    #         else:
    #             p.requires_grad = True  # 解冻 LoRA 参数

    #     if bias == 'none':
    #         return
    #     elif bias == 'all':
    #         for n, p in self.named_parameters():
    #             if 'bias' in n:
    #                 p.requires_grad = True
    #     elif bias == 'lora_only':
    #         for m in self.modules():
    #             if isinstance(m, LoRALayer) and hasattr(m, 'bias') and m.bias is not None:
    #                 m.bias.requires_grad = True
    #     else:
    #         raise NotImplementedError(f"未知 bias 选项: {bias}")
    def forward(self, x, return_features=False):
        # features = []
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        out_feature = self.conv_1(x)

        out_b1, out_b1_1, out_b1_2, out_b1_3 = self.block_1(out_feature)
        out_b2, out_b2_1, out_b2_2, out_b2_3 = self.block_2(out_b1)
        out_b3, out_b3_1, out_b3_2, out_b3_3 = self.block_3(out_b2)

        out_b4, _, _, _ = self.block_4(out_b3)
        out_b5, _, _, _ = self.block_5(out_b4)
        out_b6, out_b5_2, _, _ = self.block_6(out_b5)

        out_b6 = self.conv_2(out_b6)
        out = self.conv_cat(torch.cat([out_feature, out_b6, out_b1, out_b5_2], 1))
        output = self.upsampler(out)
        
        # features.append(out_b1_1)
        # features.append(out_b1_2)
        # features.append(out_b1_3)
        # features.append(out_b2_1)
        # features.append(out_b2_2)
        # features.append(out_b2_3)
        # features.append(out_b3_1)
        # features.append(out_b3_2)
        # features.append(out_b3_3)
 

        if return_features:
            return output, features  # Return output and intermediate features
        return output