dtype 与设备
主线位置:张量主线的第一步:张量被创建时定下 dtype 和设备,之后每一步计算都按这两个值走。
版本:PyTorch 2.15.0(CPU,设备一节另用 MPS)、burn 0.22.1(Flex 设备,对照用 ndarray 设备)
flowchart TB pg["torch 全局默认"] --> pt["float32 int64"] dev["Device 设置"] --> fx["Flex F32 I32"] dev --> nd["ndarray F32 I64"] ex["显式 dtype"] -->|优先| t["新张量"] fx --> t nd --> t t -->|cast| t2["换 dtype"] t -->|to_device| t3["换设备 dtype 不变"] t -->|"和别的 dtype 运算"| boom["后端 panic"]
本页的对照:PyTorch 侧 tensora_dtype.py,burn 侧实验 crate tensora-t 的 dtype(feature 同上一页,开着 flex 和 ndarray;debug 构建,另注明 release 的地方除外)。设备设置是进程级的全局状态,所以每个用例单独起一个进程。平台 M2 Max、macOS 27、rustc 1.99.0。panic 信息里 cargo 缓存的路径前缀用 ... 略去。
ndarray 后端在 0.22 标为弃用,这里只拿它当「另一个设备」做对照。dtype 怎么从设备设置里取出来、I32 求和溢出、Metal 不支持 F64,burn 卷的同一段代码,Int 在不同设备上位宽不同已经讲过,本页不重复。
默认 dtype:全局的,还是设备的
torch.get_default_dtype()
torch.tensor([1.0, 2.0]).dtype; torch.tensor([1, 2]).dtype; torch.arange(3).dtype
torch.set_default_dtype(torch.float64)
torch.tensor([1.0]).dtype; torch.zeros(2).dtype; torch.tensor([1]).dtype
for (name, device) in [("flex", Device::flex()), ("ndarray", Device::ndarray())] {
let s = device.settings();
let f = Tensor::<1>::from_floats([1.0, 2.0], &device);
let i = Tensor::<1, Int>::from_ints([1, 2], &device);
let r = Tensor::<1, Int>::arange(0..3, &device);
let fi = f.clone().int();
let back = i.clone().float();
// 打印 settings 和各自的 dtype
}
default dtype torch.float32
tensor([1.0, 2.0]) torch.float32 | tensor([1, 2]) torch.int64 | arange(3) torch.int64
after set_default_dtype(float64): tensor([1.0]) torch.float64 zeros torch.float64 tensor([1]) torch.int64
flex settings DeviceSettings { float_dtype: F32, int_dtype: I32, bool_dtype: Native, quantization: QuantConfig { ... } } | from_floats F32 from_ints I32 arange I32 | float.int() I32 int.float() F32
ndarray settings DeviceSettings { float_dtype: F32, int_dtype: I64, bool_dtype: Native, quantization: QuantConfig { ... } } | from_floats F32 from_ints I64 arange I64 | float.int() I64 int.float() F32
(quantization 字段很长,用 { ... } 略去。)
- PyTorch 的默认浮点类型是一个进程级的全局值,
set_default_dtype只管浮点,整数一直是 int64。 - burn 的默认值挂在设备上:每个后端设备的
defaults()给出一组 float / int / bool 类型(Flex 是F32/I32,ndarray 是F32/I64)。同一段代码换一个设备,Tensor<1, Int>里装的整数位宽就变了。 .int()转成的也是设备的默认整数类型:Flex 上I32,ndarray 上I64。PyTorch 的.int()永远是 int32,.long()才是 int64。
要改默认值,用 Device::configure,它对应的不是 torch.set_default_dtype 那样随时可改的全局开关:设置按设备登记、只能初始化一次,第一次在这个设备上建张量也算初始化。浮点和整数要一起改,就在一次调用里传一个元组:
let mut device = Device::flex();
println!("configure F64: {:?}", device.configure(FloatDType::F64));
println!("configure I64: {:?}", device.configure(IntDType::I64));
// ...
let other = Device::flex();
println!("another Device::flex() settings {:?}", other.settings());
configure F64: Ok(())
configure I64: Err(AlreadyInitialized { device: "Flex(Cpu)" })
after configure: zeros F64 arange I32
...
another Device::flex() settings DeviceSettings { float_dtype: F64, int_dtype: I32, bool_dtype: Native, quantization: QuantConfig { ... } }
let mut device = Device::flex();
println!("configure (F64, I64): {:?}", device.configure((FloatDType::F64, IntDType::I64)));
configure (F64, I64): Ok(())
zeros F64 arange I64
第一次 configure 已经把整组设置定下来(没给的整数类型取默认 I32),第二次就是 AlreadyInitialized。设置按设备登记在全局表里,不存在 Device 这个值里:另一个 Device::flex() 读到的也是 F64。configure 的文档写的是「Settings can only be initialized once per device」,怎么锁定见 burn 卷 Device 的「设置与 configure」。
问: 设备已经 configure 成 F64,zeros 也确实是 F64,为什么 Tensor::<1>::from_floats([0.1, 0.2], &device) 读回来是 0.10000000149011612?
答: 因为 from_floats 的实现是 Self::from_data(floats.into().convert::<f32>(), device):数据先被转成 f32,再按设备默认值升回 F64,精度在第一步就丢了。同一个设备上用 from_data,f64 的数据原样进来:
from_floats on F64 device -> F64 values [0.10000000149011612, 0.20000000298023224]
from_data f64 on F64 device -> F64 values [0.1, 0.2]
要 F64 的精度,用 from_data,不要用 from_floats。
numpy 来的 f64
PyTorch 从 numpy 拿数据时保留 dtype:np.array([0.1, 0.2]) 是 float64,torch.from_numpy 和 torch.tensor 都给出 float64。burn 的 from_data 不看数据的元素类型,只看设备默认值,f64 的数据会被静静地转成 F32:
n = np.array([0.1, 0.2])
torch.from_numpy(n).dtype, torch.tensor(n).dtype
torch.tensor(n, dtype=torch.float32).double().tolist()
let a = Tensor::<1>::from_data([0.1f64, 0.2], &device);
let b = Tensor::<1>::from_data([0.1f64, 0.2], (&device, DType::F64));
let big = Tensor::<1, Int>::from_data([1i64 << 40], (&device, DType::I64));
from_numpy f64 torch.float64 | tensor(np f64) torch.float64
tensor(np f64, dtype=float32) values [0.10000000149011612, 0.20000000298023224]
from_data f64 -> F32 values [0.10000000149011612, 0.20000000298023224]
from_data f64 + DType::F64 -> F64 values [0.1, 0.2]
Int with DType::I64 -> I64 Ok([1099511627776])
没写 dtype 时 burn 的结果和 PyTorch 里显式写 dtype=torch.float32 一样,两边读回的值逐位相同。要保留 f64,把 (&device, DType::F64) 当第二个参数传进去:元组也实现了 Into<TensorCreationOptions>,from_data 里的 opt.resolve_dtype::<K>() 先用显式给的 dtype,没有才取设备默认值。
转换:.to(dtype) 对 cast
x = torch.tensor([1.5, -2.7])
x.int(); x.long(); x.to(torch.float64); x.half()
let a = Tensor::<1>::from_floats([1.5, -2.7], &device);
let a64 = a.clone().cast(FloatDType::F64);
let h = a.clone().cast(FloatDType::F16);
let i: Tensor<1, Int> = a.clone().cast(IntDType::I64);
let j = a.clone().int();
.int() torch.int32 [1, -2] | .long() torch.int64 | .to(float64) torch.float64 | .half() [1.5, -2.69921875]
cast F64 -> F64
cast F16 -> F16 [1.5, -2.69921875]
cast IntDType::I64 -> Tensor<1, Int> I64 Ok([1, -2])
int() -> I32 Ok([1, -2])
数值上两边一致:浮点转整数都向零截断,F16 的 -2.7 都读回 -2.69921875。不同在类型上:cast 的返回类型是 Tensor<D, T::OutputKind>,Cast trait 按参数的类型决定输出种类:传 FloatDType 还是 Tensor<D>,传 IntDType 就变成 Tensor<D, Int>。PyTorch 的 .to(torch.int64) 只改一个运行期的属性,burn 里种类变了,变量的类型也跟着变。
参数写成泛泛的 DType 时,种类对不上只能到运行期才发现。Tensor<1> 上 cast(DType::I32):
thread 'main' (8547784) panicked at .../burn-std-0.22.1/src/tensor/dtype.rs:215:18:
Expected float data type, got I32
没有类型提升
PyTorch 在二元运算里自动提升类型(type promotion):
x64 = x.double()
(x64 + x).dtype, (x + x64).dtype
(torch.tensor([1, 2], dtype=torch.int32) + torch.tensor([1, 2])).dtype
(torch.tensor([1, 2]) + torch.tensor([0.5, 0.5])).dtype
(torch.tensor([1, 2]) + 0.5).dtype
f64 + f32 -> torch.float64 | f32 + f64 -> torch.float64
i32 + i64 -> torch.int64
i64 + f32 -> torch.float32
i64 tensor + 0.5 -> torch.float32
burn 没有这一层。同种类、不同 dtype 的两个张量相加,在后端里 panic:
let a = Tensor::<1>::from_floats([1.5, -2.7], &device);
let a64 = a.clone().cast(FloatDType::F64);
let sum = a64 + a;
### mix-float flex
thread 'main' (8547820) panicked at .../burn-flex-0.22.1/src/ops/binary.rs:41:5:
assertion `left == right` failed: binary_op: dtype mismatch
left: F64
right: F32
### mix-float ndarray
thread 'main' (8547824) panicked at .../burn-ndarray-0.22.1/src/ops/tensor.rs:87:9:
Data type mismatch (lhs: F64, rhs: F32)
### mix-int flex
thread 'main' (8547828) panicked at .../burn-flex-0.22.1/src/ops/binary.rs:361:5:
assertion `left == right` failed: int_binary_op: dtype mismatch
left: I32
right: I64
### release mix-float flex
thread 'main' (8547836) panicked at .../burn-flex-0.22.1/src/tensor.rs:154:9:
storage: dtype mismatch (expected F32, got F64)
(每段只留 panic 行和原因,note: run with RUST_BACKTRACE=1 行略去。)
这些 panic 都不是 TensorCheck 报的。add 的入口检查 binary_ops_ew 只查设备和形状,不查 dtype;dtype 由各个后端自己把关。Flex 的 binary_op 第一行是 debug_assert_eq!(lhs.dtype(), rhs.dtype(), "binary_op: dtype mismatch"),release 下这行不存在,panic 改由读存储的 FlexTensor::storage 里的 assert! 抛出。所以同一个错误,debug 和 release 的报错位置、措辞都不一样。release 那条的措辞还是反的:格式串是 "(expected {:?}, got {:?})", self.dtype, E::dtype(),「expected」后面打印的其实是张量自己的 dtype(右操作数 F32),「got」是代码想按哪种类型读(左操作数决定的 F64)。
种类不同(Tensor<1, Int> + Tensor<1>)则是编译错误,报错说 Tensor<1>: ElementConversion 没实现,读法见归约与广播里广播的那一节。
最容易漏掉的是整数张量和浮点标量:
let i = Tensor::<1, Int>::from_ints([1, 2], &device);
i.clone() + 0.5;
i * 1.5;
int + 0.5 [2] float32 [1.5, 2.5]
int * 1.5 [2] float32 [1.5, 3.0]
Int + 0.5 [2] I32 [1, 2]
Int * 1.5 [2] I32 [1, 2]
(这四行输出来自下一页的对照程序 tensora_reduce.py 和 reduce。)没有报错,结果是错的。Tensor + 标量 走的是 add_scalar,它先用 Scalar::new 把标量换成张量的种类:张量是整数,就是 Self::Int(value.elem()),0.5 截断成 0、1.5 截断成 1。输出的种类永远是左边张量的种类。要浮点结果,先 i.float() 再算。
设备:运行期的值,这点两边像
torch.device 是运行期的值,x.to("mps") 换设备。burn 0.22 去掉了后端类型参数之后,Device 也是一个普通的值,可以在运行期挑:
print("x.device", x.device, "| mps available", torch.backends.mps.is_available())
m = x.to("mps")
x.reshape(2, 1) + m.reshape(2, 1).expand(2, 3)
x.double().to("mps")
let use_ndarray = std::env::var("USE_NDARRAY").is_ok();
let device = if use_ndarray { nd.clone() } else { flex.clone() };
println!("picked at runtime: {device:?}");
// ...
let a = Tensor::<1>::from_floats([1.0, 2.0], &Device::flex());
let b = Tensor::<1>::from_floats([1.0, 2.0], &Device::ndarray());
let c = a + b;
x.device cpu | mps available True
to mps mps:0 torch.float32
cpu + mps: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!
mps float64: TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
picked at runtime: Device<Flex(Cpu)>
picked at runtime: Device<NdArray(Cpu)>
(第二行是设了 USE_NDARRAY=1 的运行。)
thread 'main' (8547816) panicked at .../burn-tensor-0.22.1/src/tensor/api/numeric.rs:40:9:
=== Tensor Operation Error ===
Operation: 'Add'
Reason:
1. The provided tensors are not on the same device. Lhs tensor device Device<Flex(Cpu)>, Rhs tensor device Device<NdArray(Cpu)>.
两边都在运算时检查设备、都不自动搬运。burn 的设备检查在 TensorCheck 里,所以和形状检查一样,release 下也开着。Device::default() 相当于 PyTorch 里「没写 device」:它按编译进来的后端挑一个(本页的 feature 下是 Device<Flex(Cpu)>),也可以用环境变量 BURN_DEVICE 指定,BURN_DEVICE=ndarray 时打印 Device<NdArray(Cpu)>;挑选顺序见 Device。MPS 拒绝 float64 的那条,在 burn 里对应 Metal 设备不支持 F64,见 dtype-per-device。
to_device 对应 .to(device),只换设备,不换 dtype,和 PyTorch 的 .to("mps") 保留 float32 一样。
let flex = Device::flex();
let nd = Device::ndarray();
let moved = Tensor::<1, Int>::arange(0..3, &flex).to_device(&nd);
let fresh = Tensor::<1, Int>::arange(0..3, &nd);
let sum = moved + fresh;
动脑筋答案
Int on flex I32 -> to_device(ndarray) I32 on Device<NdArray(Cpu)>
fresh Int on ndarray I64
thread 'main' (8547812) panicked at .../burn-ndarray-0.22.1/src/ops/int_tensor.rs:152:9:
Data type mismatch (lhs: I32, rhs: I64)
panic。设备检查通过了(两个都在 ndarray 上),但 moved 从 Flex 带来了 I32,fresh 是 ndarray 默认的 I64。上一节说过 burn 没有类型提升,ndarray 后端看到两个 dtype 不同就 panic。dtype 在创建时由当时的设备定下,之后跟着张量走,to_device 不会把它换成目标设备的默认值。混用设备时,在 to_device 之后显式 cast 到你要的类型,或者创建时就写出 dtype。
小结
- PyTorch 的默认浮点 dtype 是全局的;burn 的默认 dtype 属于设备(Flex
F32/I32,ndarrayF32/I64),.int()也取设备默认值。 configure每个设备只能成功一次,浮点和整数要放在一个元组里一起改;设置登记在全局表里,同一个设备的所有Device值共享。from_data不保留数据的元素类型,f64数据默认变成F32;要精度就传(&device, DType::F64)。from_floats总是先转成f32。cast(IntDType::..)改变张量的种类,返回Tensor<D, Int>;泛泛的DType参数种类不对时运行期 panic。- 没有类型提升:同种类不同 dtype 相加由后端 panic,debug 和 release 报错不同;整数张量加浮点标量静默截断标量。
- 设备两边都是运行期的值,跨设备运算都报错;
to_device不改 dtype,混用设备后要自己cast。