使用k.variable、add_weight的TensorFlow层转PyTorch后性能低如何修正

核心错误点 #

• 非负约束逻辑错误：原TensorFlow代码仅对可训练参数it_weights施加非负约束，你的实现是将kernel * it_weights的结果整体过ReLU，会抹除DCT核本身的合法负值，导致特征变换完全失效，这是性能下降的最主要原因。
• 不必要的硬编码与拆分冗余：你硬编码了.to('cuda')，会导致设备不匹配报错；拆分了三个不必要的类，既增加了冗余也容易引入变量传递的错误。
• DCT核可训练性设置问题：原逻辑中DCT核是固定的基函数，不需要参与训练，需要显式设置禁止梯度更新避免不必要的参数变动。
• 变量覆盖隐患：原TensorFlow代码call方法直接修改self.kernel成员变量的写法存在逻辑隐患，会导致每次前向传播覆盖原始DCT核，需要避免同类问题。

修正后的PyTorch实现 #

import torch
import torch.nn as nn
import torch.nn.functional as F
from math import sqrt, cos, pi

class adaptive_implicit_trans(nn.Module):
    def __init__(self):
        super().__init__()
        # 可训练非负权重，对应原TF的it_weights
        self.it_weights = nn.Parameter(torch.ones(1, 64, 1, 1), requires_grad=True)
        # 初始化固定DCT核
        conv_shape = (64, 64, 1, 1)
        kernel = torch.zeros(conv_shape)
        r1 = sqrt(1.0 / 8)
        r2 = sqrt(2.0 / 8)
        for i in range(8):
            _u = 2 * i + 1
            for j in range(8):
                _v = 2 * j + 1
                out_idx = i * 8 + j
                for u in range(8):
                    for v in range(8):
                        in_idx = u * 8 + v
                        t = cos(_u * u * pi / 16) * cos(_v * v * pi / 16)
                        t = t * r1 if u == 0 else t * r2
                        t = t * r1 if v == 0 else t * r2
                        kernel[out_idx, in_idx, 0, 0] = t
        # 固定DCT核到buffer，不参与训练、自动跟随模型设备迁移
        self.register_buffer('kernel', kernel)
    
    def forward(self, inputs):
        # 仅对it_weights做非负约束，保留kernel的正负值
        clamped_weights = torch.clamp(self.it_weights, min=0.0)
        conv_kernel = self.kernel * clamped_weights
        # 1x1卷积padding=0等价same padding，显式声明对齐原逻辑
        y = F.conv2d(inputs, conv_kernel, padding=0)
        return y

    def compute_output_shape(self, input_shape):
        return input_shape

额外注意事项 #

• 确保输入PyTorch模型的张量格式为[batch_size, channels, height, width]（PyTorch默认的channels first格式），和TensorFlow的channels last格式区分开。
• PyTorch 0.4之后版本不需要手动调用Variable，nn.Parameter已经自动支持梯度追踪。

内容的提问来源于stack exchange，提问作者Junghyun Kim