TensorFlow中Conv2D扩展二值权重后反向传播算子适配问题咨询

问题背景 #

我正在TensorFlow框架中扩展Conv2D算子以支持二值权重，已经完成了以下几步操作：

• 修改tensorflow/core/ops/nn_ops.cc，给REGISTER_OP("Conv2D")添加了use_binary_weight布尔属性，用来标识是否启用二值权重卷积：

  REGISTER_OP("Conv2D")
      .Input("input: T")
      .Input("filter: T")
      .Output("output: T")
      .Attr("T: {half, bfloat16, float, double}")
      .Attr("strides: list(int)")
      .Attr("use_cudnn_on_gpu: bool = true")
      .Attr(GetPaddingAttrString())
      .Attr(GetConvnetDataFormatAttrString())
      .Attr("dilations: list(int) = [1, 1, 1, 1]")
      .SetShapeFn(shape_inference::Conv2DShape)
      .Attr("use_binary_weight: bool = true");
  

• 修改tensorflow/core/kernels/conv_ops.cc，给Conv2DOp类添加了私有布尔成员use_binary_weight_，并在构造函数中通过下面的代码赋值：

  OP_REQUIRES_OK(context, context->GetAttr("use_binary_weight", &use_binary_weights_));
  

• 在Conv2DOp::Compute()方法开头，当use_binary_weight_为true时，把全精度权重转换为二值权重，完成了二值卷积的前向计算。

现在遇到了一个棘手的问题：二值卷积的梯度计算逻辑和标准卷积不一样，需要对Conv2DBackpropInput和Conv2DBackpropFilter这两个算子做特殊处理，但这两个算子并不是在训练脚本里显式创建的，而是TensorFlow自动生成的梯度节点。我该怎么让这两个梯度算子知道当前是否启用了二值卷积呢？

解决方案思路 #

我之前做类似的TensorFlow算子扩展时，刚好踩过这个坑，下面几个思路都是实际验证过可行的：

1. 给梯度算子添加对应属性，在梯度注册时传递标识 #

这是最贴合TensorFlow原生设计的方案：

• 第一步：给梯度算子添加use_binary_weight属性
  修改tensorflow/core/ops/nn_ops.cc，给Conv2DBackpropInput和Conv2DBackpropFilter也加上和Conv2D一致的属性：

  // 给Conv2DBackpropInput添加属性
  REGISTER_OP("Conv2DBackpropInput")
      // 保留原有所有Input/Output/Attr定义
      .Attr("use_binary_weight: bool = true");
  
  // 给Conv2DBackpropFilter添加属性
  REGISTER_OP("Conv2DBackpropFilter")
      // 保留原有所有Input/Output/Attr定义
      .Attr("use_binary_weight: bool = true");
  

• 第二步：修改Conv2D的梯度注册逻辑，传递属性值
  找到tensorflow/core/ops/nn_ops.cc中Conv2D的梯度注册代码（一般是REGISTER_GRADIENT_OP绑定的梯度函数），在生成梯度节点时，把原Conv2D节点的use_binary_weight属性传递给梯度算子：

  REGISTER_GRADIENT_OP("Conv2D", Conv2DGrad);
  
  Status Conv2DGrad(const GradientTensor& grad_inputs,
                    const GradientTensor& forward_inputs,
                    std::vector<GradientTensor>* grad_outputs) {
    // 先获取原Conv2D节点的use_binary_weight属性值
    bool use_binary_weight;
    TF_RETURN_IF_ERROR(grad_inputs.op->GetAttr("use_binary_weight", &use_binary_weight));
  
    // 构建Conv2DBackpropInput节点时，传入属性
    NodeDefBuilder backprop_input_builder(strings::StrCat(grad_inputs.op->name(), "_backprop_input"), 
                                         "Conv2DBackpropInput");
    // 设置原有输入参数（比如input_sizes, filter, out_backprop等）
    // ...
    // 添加二值权重标识属性
    backprop_input_builder.Attr("use_binary_weight", use_binary_weight);
    // 完成节点构建并添加到梯度输出
    // ...
  
    // 同样处理Conv2DBackpropFilter节点
    NodeDefBuilder backprop_filter_builder(strings::StrCat(grad_inputs.op->name(), "_backprop_filter"), 
                                         "Conv2DBackpropFilter");
    // 设置原有输入参数（比如input, filter_sizes, out_backprop等）
    // ...
    backprop_filter_builder.Attr("use_binary_weight", use_binary_weight);
    // 完成节点构建并添加到梯度输出
    // ...
  
    return Status::OK();
  }
  

• 第三步：修改梯度算子的Kernel实现
  和你处理Conv2DOp的方式一样，在Conv2DBackpropInputOp和Conv2DBackpropFilterOp类中添加use_binary_weight_成员，在构造函数中读取属性值，然后在Compute方法中根据这个标识执行二值卷积对应的梯度计算逻辑。

2. 注册自定义梯度函数，替换默认梯度实现 #

如果不想修改太多原生梯度算子的代码，你可以给带use_binary_weight=true的Conv2D注册一个完全自定义的梯度函数，直接调用你自己实现的二值反向传播算子：

Status Conv2DGrad(const GradientTensor& grad_inputs,
                  const GradientTensor& forward_inputs,
                  std::vector<GradientTensor>* grad_outputs) {
  bool use_binary_weight;
  TF_RETURN_IF_ERROR(grad_inputs.op->GetAttr("use_binary_weight", &use_binary_weight));

  if (use_binary_weight) {
    // 这里直接调用你自己实现的二值卷积反向传播算子
    // 比如创建自定义的BinaryConv2DBackpropInput和BinaryConv2DBackpropFilter节点
    NodeDefBuilder backprop_input_builder(/*节点名称*/, "BinaryConv2DBackpropInput");
    // 设置输入参数、属性（比如二值权重相关配置）
    // ...
    // 构建节点并加入梯度输出
  } else {
    // 执行原生Conv2D的梯度计算逻辑，和原有代码保持一致
    // ...
  }

  return Status::OK();
}

这种方式的好处是把二值卷积的前后向逻辑完全和原生逻辑隔离开，避免互相干扰，适合需要做大量定制化梯度逻辑的场景。

3. 临时方案：通过张量自定义属性传递标识（不推荐） #

如果上面两种方式都需要修改太多代码，你可以尝试在Conv2DOp::Compute中给权重张量附加一个自定义属性，比如：

// 在Conv2DOp::Compute中，当use_binary_weight_为true时
filter_tensor->SetAttr("is_binary_weight", true);

然后在梯度算子的Compute方法中读取这个属性：

bool is_binary_weight = false;
filter_tensor->GetAttr("is_binary_weight", &is_binary_weight);
if (is_binary_weight) {
  // 执行二值梯度逻辑
}

不过这种方式不符合TensorFlow的设计规范，因为张量属性应该用来描述张量本身的特性，而不是算子的行为，只建议作为临时应急方案使用。

内容的提问来源于stack exchange，提问作者Charles Tao