Rust in Everything
基于 burn 0.22.1

Int + 0.5 还是 Int,小数被截掉

主线位置:张量主线的「计算」一步,整数张量和浮点标量的二元运算。
版本:PyTorch 2.15.0(CPU)、burn 0.22.1(Flex 设备)
flowchart TB
  idx["idx Int 0 到 3"] --> pt["PyTorch idx + 0.5"]
  idx --> bn["burn idx + 0.5"]
  pt --> f["float32 0.5 到 3.5"]
  bn --> cut["0.5 先截成 0"]
  cut --> same["I32 0 到 3 不变"]
  same --> fix["idx.float() + 0.5"]

实验 crate pita-t 的 p03、编不过的例子在 pita-ce 的 ce03(同样的 burn feature),PyTorch 侧 p03_int_plus_float.py。debug 构建。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

像素中心、归一化坐标、位置编码,起点都是一个整数下标张量:

idx = torch.arange(4)            # pixel indices, int64
centers = (idx + 0.5) / 4        # normalized pixel centers
idx + 0.5       [0.5, 1.5, 2.5, 3.5] torch.float32
(idx+0.5)/4     [0.125, 0.375, 0.625, 0.875]
idx * 0.5       [0.0, 0.5, 1.0, 1.5]
idx * -0.7      [-0.0, -0.699999988079071, -1.399999976158142, -2.0999999046325684]

整数张量遇到浮点标量,结果提升成默认浮点类型 float32。

搬到 burn 的现象

let idx = Tensor::<1, Int>::arange(0..4, &device);
let shifted = idx.clone() + 0.5;
println!("(idx + 0.5) / 4 {:?}", iv(shifted / 4));
println!("idx * 0.5       {:?}", iv(idx.clone() * 0.5));
println!("idx * -0.7      {:?}", iv(idx.clone() * -0.7));
println!("idx * 2.9       {:?}", iv(idx.clone() * 2.9));
idx + 0.5       I32 [0, 1, 2, 3]
(idx + 0.5) / 4 [0, 0, 0, 0]
idx * 0.5       [0, 0, 0, 0]
idx * -0.7      [0, 0, 0, 0]
idx * 2.9       [0, 2, 4, 6]

能编译,能运行,没有警告。idx + 0.5 原样返回 idx,乘 0.5 和 -0.7 得到全 0,乘 2.9 等于乘 2。结果的种类还是 Int,dtype 是 I32。

原因

idx + 0.5 用的是 Add<E: ElementConversion> 这一个重载,转给 add_scalar,乘法同理走 mul_scalar。它们先用 Scalar::new 把标量换成「张量自己的种类」:

    pub fn new<E: ElementConversion>(value: E, dtype: &DType) -> Self {
        if dtype.is_float() | matches!(dtype, &DType::QFloat(_)) {
            Self::Float(value.elem())
        } else if dtype.is_int() {
            Self::Int(value.elem())
        // ...

张量是整数,0.5 就经 value.elem() 变成 i64。浮点转整数的规则在 burn-std 的 impl_to_element_float_to_signed_int 里,注释写着「Float as int truncates toward zero」:0.5、-0.7 变成 0,2.9 变成 2。只有超出 i64 的范围才 panic:idx + 1e30 报「Float cannot be represented in the target signed int type」。运算的输出种类永远是左边张量的种类,类型提升(type promotion)这一层在 burn 里不存在,见dtype 与设备的「没有类型提升」。

改法

先把整数张量换成浮点,再和浮点标量运算:

let centers = (idx.float() + 0.5) / 4.0;
(idx.float() + 0.5) / 4.0 F32 [0.125, 0.375, 0.625, 0.875]

和 PyTorch 的 [0.125, 0.375, 0.625, 0.875] 逐个相同:这几个数在 f32 里都能精确表示,两边算的是同样的 f32 运算,误差是 0。.float() 取设备的默认浮点类型(Flex 上 F32),要别的精度用 cast(FloatDType::F64)。

怎么检查

结果的种类写在类型里,所以给结果标上你以为的类型,编译器就会拦住:

let centers: Tensor<1> = (idx + 0.5) / 4.0;
error[E0308]: mismatched types
 --> src/bin/ce03.rs:6:28
  |
6 |   let centers: Tensor<1> = (idx + 0.5) / 4.0;
  |                ---------   ^^^^^^^^^^^^^^^^^ expected `Tensor<1>`, found `Tensor<1, Int>`
  |                |
  |                expected due to this
  |
  = note: expected struct `Tensor<_, burn::tensor::Float>`
             found struct `Tensor<_, burn::tensor::Int>`
  • 函数签名里写明 Tensor<D>(浮点)还是 Tensor<D, Int>,不要让中间结果的类型全靠推断。
  • 在搬过来的代码里找 Int 张量和带小数点的字面量放在同一个运算里的地方:+ 0.5、* 0.1、/ 2.0、- 1e-3。每一处都要么先 .float(),要么确认截断就是你要的。

小结

  • burn 没有类型提升:Int 张量和浮点标量运算,标量先向零截成整数,结果还是 Int。
  • idx + 0.5 不变、idx * 0.5 全 0、idx * 2.9 等于乘 2,都不报错。
  • 先 .float() 再算;给结果标上 Tensor<D> 让编译器检查种类。

💬 关联讨论

发起讨论
加载中...