ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

PyTorch张量类型转换实战:从dtype报错到混合精度训练

PyTorch张量类型转换实战:从dtype报错到混合精度训练 1. 张量类型转换到底在解决什么问题刚接触深度学习框架的人十有八九会在某个深夜被一行报错拦住去路RuntimeError: expected scalar type Float but found Double或者TypeError: expected TensorOptions(dtypefloat, ...)。这类报错翻来覆去就一个根源——张量的数据类型不匹配。张量类型转换这件事说白了就是把一个张量从一种数据类型变成另一种比如把float64变成float32把int64变成float32或者把 GPU 上的张量搬到 CPU 上同时换类型。听起来简单但真正踩过坑的人都知道这里面藏着精度丢失、显存翻倍、梯度断裂、设备不匹配等一堆细节。我自己最早做模型训练的时候数据预处理用 NumPy 读进来默认是float64直接转成张量喂给网络结果第一层卷积就报错。当时不懂以为是网络定义有问题查了半天才发现是 dtype 没统一。后来做混合精度训练又遇到float16和float32之间来回转换导致 loss 变成 NaN 的情况。这些经历让我意识到张量类型转换不是一个“知道有这回事就行”的知识点而是每天写代码都会碰到的基本功。这篇文章适合谁看如果你正在学 PyTorch、TensorFlow 或者任何涉及张量操作的框架如果你被 dtype 报错折磨过如果你想搞清楚float32、float64、float16、bfloat16、int8这些类型到底该怎么选、怎么转、什么时候转那这篇内容就是给你写的。我会从底层逻辑讲到实操代码从常见报错讲到排查技巧尽量把这件事讲透。2. 张量数据类型体系与转换的底层逻辑2.1 张量有哪些数据类型各自什么场景用不同框架的类型命名略有差异但核心体系是相通的。以 PyTorch 为例常用的浮点类型有torch.float32也叫torch.float、torch.float64也叫torch.double、torch.float16也叫torch.half、torch.bfloat16。整数类型有torch.int8、torch.int16、torch.int32、torch.int64也叫torch.long还有无符号的torch.uint8。布尔类型是torch.bool复数类型有torch.complex64和torch.complex128。这些类型不是随便设计的每一种都有明确的用途。float32是深度学习的默认类型单精度浮点数占 4 个字节精度和显存占用之间平衡得最好。float64是双精度占 8 个字节精度更高但显存翻倍通常只在科学计算或者需要高精度累加的场景用。float16是半精度占 2 个字节主要用于混合精度训练能省显存、加速计算但精度低容易溢出。bfloat16也是 2 个字节但指数位和float32一样多动态范围更大在 A100 之后的显卡上表现很好。整数类型里int64是最常用的索引、标签、embedding 的输入都用它。int32在某些场景下能省内存但很多框架的默认索引类型是int64混用容易出问题。uint8常用于图像数据像素值 0 到 255 正好对应一个字节。bool用于掩码操作比如 attention mask。提示PyTorch 的默认浮点类型是float32默认整数类型是int64。当你用torch.tensor([1, 2, 3])创建张量时得到的是int64用torch.tensor([1.0, 2.0])创建时得到的是float32。这个默认行为可以通过torch.set_default_dtype()修改但一般不建议改容易让代码变得难以预测。2.2 类型转换的本质重新解释内存中的比特位理解类型转换的底层关键在于明白一件事转换不是简单地换个标签而是重新解释甚至重新计算内存中的比特位。对于整数之间的转换比如int32转int64底层是把 4 个字节扩展成 8 个字节数值不变。但反过来int64转int32如果数值超过int32的范围约 21 亿就会溢出结果不可预测。浮点之间的转换更复杂float32转float16时指数位从 8 位缩到 5 位尾数从 23 位缩到 10 位超出范围的值会变成inf精度不够的值会被截断。浮点和整数之间的转换也有讲究。float32转int64时小数部分直接被截断不是四舍五入。比如3.9转成整数是3-3.9转成整数是-3。如果你需要四舍五入得先用torch.round()。反过来int64转float32大整数可能丢失精度因为float32的尾数只有 23 位能精确表示的整数上限是 2 的 24 次方也就是 16777216。超过这个值转换结果就不精确了。import torch # 浮点转整数截断而非四舍五入 a torch.tensor([3.9, -3.9, 3.1, -3.1]) print(a.to(torch.int64)) # tensor([ 3, -3, 3, -3]) # 大整数转 float32 丢精度 b torch.tensor([16777217], dtypetorch.int64) print(b.to(torch.float32)) # tensor([16777216.])这段代码我第一次跑的时候挺惊讶的16777217转成float32居然变成了16777216。后来查了 IEEE 754 的规范才明白float32的尾数位只有 23 位加上隐含的 1 位总共能精确表示 24 位二进制整数也就是 2 的 24 次方。超过这个范围就只能近似表示了。2.3 为什么框架要设计这么多类型统一用 float32 不行吗有人可能会想既然float32这么好用为什么还要搞这么多类型出来全部用float32不就完了这个问题我当初也问过自己后来在实际项目里才慢慢理解。首先是显存。一个 10 亿参数的模型用float32存权重需要 4GB 显存用float16只需要 2GB。训练的时候还有梯度、优化器状态、激活值显存占用是参数量的好几倍。混合精度训练能把显存占用降下来靠的就是float16和float32的配合。其次是计算速度。现代 GPU 对float16和bfloat16有专门的加速单元矩阵乘法的吞吐量能比float32高好几倍。在 A100 上bfloat16的矩阵乘法性能是float32的 8 倍左右。这不是小数目训练大模型的时候能省下大量时间。再就是精度需求。有些场景就是需要float64比如数值分析、科学计算、某些强化学习的价值估计。用float32会导致累积误差过大结果不可信。而有些场景只需要int8比如量化推理模型权重和激活值都压到 8 位整数推理速度能提升好几倍精度损失在可接受范围内。所以类型转换不是可有可无的操作而是根据场景在精度、速度、显存之间做权衡的手段。3. 主流框架中张量类型转换的实操方法3.1 PyTorch 中的类型转换to、type、float 等方法全解析PyTorch 提供了好几种类型转换的方式新手容易搞混我一个个说。最通用的是.to()方法。它可以指定 dtype也可以指定 device还可以同时指定。写法是tensor.to(dtype)或者tensor.to(device)或者tensor.to(device, dtype)。比如x.to(torch.float32)、x.to(cuda)、x.to(cuda, torch.float16)。.to()的好处是灵活一个方法搞定设备和类型。如果目标类型和当前类型一致它会直接返回原张量不会复制所以不用担心性能问题。import torch x torch.tensor([1, 2, 3]) # int64 y x.to(torch.float32) # 转 float32 z x.to(cuda, torch.float16) # 转 float16 并搬到 GPU.type()方法是老式的写法接受字符串或者 dtype 对象。比如x.type(torch.FloatTensor)或者x.type(torch.FloatTensor)。这种写法在旧代码里常见现在更推荐用.to()。.type()的一个特点是它只能转类型不能转设备要转设备还得配合.cuda()或.cpu()。.float()、.double()、.half()、.long()、.int()、.bool()这些是快捷方法分别对应转成float32、float64、float16、int64、int32、bool。写起来简洁但只能转类型不能转设备。我平时写代码如果只是单纯转类型用这些快捷方法如果要同时转设备和类型用.to()。x torch.tensor([1, 2, 3]) print(x.float()) # tensor([1., 2., 3.]) print(x.double()) # tensor([1., 2., 3.], dtypetorch.float64) print(x.half()) # tensor([1., 2., 3.], dtypetorch.float16) print(x.bool()) # tensor([True, True, True])还有一个容易忽略的点.to()返回的是新张量原张量不变。但如果是原地操作可以用.to_()或者.float_()这类带下划线的方法直接修改原张量。不过原地操作要小心如果原张量被其他地方引用改了会出问题。注意tensor.long()和tensor.to(torch.int64)是等价的但tensor.int()转的是int32不是int64。这个坑我踩过用int()转出来的索引张量在某些操作里会报类型错误因为框架期望的是int64。索引相关的场景一律用long()或者.to(torch.int64)。3.2 TensorFlow 中的类型转换tf.cast 与 tf.convert_to_tensorTensorFlow 的类型转换主要靠tf.cast()。写法是tf.cast(x, dtype)比如tf.cast(x, tf.float32)。它支持张量、NumPy 数组、Python 列表等各种输入返回指定类型的张量。import tensorflow as tf x tf.constant([1, 2, 3], dtypetf.int64) y tf.cast(x, tf.float32) print(y) # tf.Tensor([1. 2. 3.], shape(3,), dtypefloat32)tf.convert_to_tensor()是另一个常用函数它能把各种输入转成张量同时可以指定 dtype。如果输入已经是张量它会检查类型是否匹配不匹配就报错或者转换取决于参数。这个函数在写自定义层或者数据管道的时候很常用。x [1, 2, 3] y tf.convert_to_tensor(x, dtypetf.float32) print(y) # tf.Tensor([1. 2. 3.], shape(3,), dtypefloat32)TensorFlow 的一个特点是它对类型要求很严格。两个张量做运算如果 dtype 不一致直接报错不会自动转换。这一点和 PyTorch 类似但 TensorFlow 的报错信息有时候更晦涩需要仔细看。3.3 NumPy 与张量之间的类型转换桥接的注意事项实际项目里数据往往先用 NumPy 处理再转成张量喂给模型。这个桥接过程是类型问题的高发区。NumPy 的默认浮点类型是float64默认整数类型是int64在 64 位系统上。而 PyTorch 的默认浮点类型是float32。所以torch.from_numpy(np_array)得到的张量如果np_array是float64张量也是float64。直接喂给模型就会报类型错误。import numpy as np import torch np_array np.array([1.0, 2.0, 3.0]) # float64 tensor torch.from_numpy(np_array) print(tensor.dtype) # torch.float64 # 正确做法先转 float32 再转张量 tensor torch.from_numpy(np_array.astype(np.float32)) print(tensor.dtype) # torch.float32TensorFlow 这边tf.convert_to_tensor(np_array)会保留 NumPy 的 dtype。如果np_array是float64得到的张量也是float64。所以同样需要先astype。提示torch.from_numpy()创建的张量和原 NumPy 数组共享内存。如果你改了张量NumPy 数组也会变反之亦然。这个特性有时候很有用能省内存但如果不注意会导致难以排查的 bug。如果需要独立的内存用torch.tensor(np_array)或者torch.from_numpy(np_array).clone()。3.4 设备与类型同时转换GPU/CPU 场景下的最佳实践做深度学习张量在 CPU 和 GPU 之间搬来搬去是家常便饭。.to()方法可以同时指定设备和类型这是最推荐的做法。device torch.device(cuda if torch.cuda.is_available() else cpu) x torch.tensor([1.0, 2.0, 3.0]) x x.to(device, torch.float16)这里有个细节.to()的参数顺序不影响结果x.to(device, dtype)和x.to(dtype).to(device)效果一样。但性能上一次性.to(device, dtype)比分开两次转要快因为它只做一次内存拷贝。还有一个常见的坑从 GPU 转回 CPU 的时候如果张量需要梯度.cpu()之后梯度还在但.detach().cpu()会断开梯度。如果你只是想把 loss 拿出来打印用.item()更合适它直接返回 Python 标量。loss criterion(output, target) print(loss.item()) # 推荐直接取标量 # print(loss.cpu().numpy()) # 不推荐多此一举还可能带梯度4. 类型转换中的精度陷阱与性能考量4.1 float32 转 float16什么时候会溢出怎么防float16的动态范围比float32小很多。float32能表示的最大值约 3.4e38float16只能表示到 65504。超过这个值就变成inf。训练的时候loss 或者梯度很容易超过这个范围导致 NaN。我做过一个实验把一个简单的全连接网络的权重初始化成float16前向传播没问题但反向传播的时候梯度累加几次就溢出了。后来改成混合精度训练用torch.cuda.amp自动管理才稳定下来。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()autocast会自动把适合的操作用float16算不适合的保持float32。GradScaler会把 loss 放大防止梯度下溢更新的时候再缩回来。这套机制是混合精度训练的标准做法比自己手动转类型靠谱得多。如果不用amp手动转float16的时候一定要检查数值范围。可以用torch.isinf()和torch.isnan()检查发现异常就回退到float32。4.2 float64 转 float32精度损失有多大值不值得float64转float32尾数从 52 位缩到 23 位精度损失大约是 1e-7 到 1e-8 的量级。对于大多数深度学习任务这个精度足够了。但有些场景不行比如计算损失函数的时候如果 loss 本身很小1e-10 量级转成float32可能直接变成 0。我遇到过一个案例做物理模拟的神经网络输入输出的数值范围跨越十几个数量级用float32训练出来的模型误差很大。后来把关键计算改成float64误差才降下来。所以要不要转取决于你的数值范围和对精度的要求。一个折中的做法是数据存储和传输用float32关键累加用float64。比如计算均值、方差的时候先转float64算完再转回来。这样既省内存又保证精度。x torch.randn(1000000, dtypetorch.float32) mean x.to(torch.float64).mean().to(torch.float32)4.3 int64 与 int32 的取舍索引、标签、embedding 的类型选择int64和int32的选择主要看用途。索引和标签一般用int64因为框架的默认索引类型是int64混用会报错。Embedding 的输入必须是int64这是硬性要求。int32的优势是省内存。一个 100 万长度的标签数组int64占 8MBint32占 4MB。数据量大的时候这个差距会累积。但省下来的内存往往不值得带来的类型转换麻烦。我的建议是除非明确知道内存是瓶颈否则统一用int64省心。有一个例外是图像数据。图像的像素值范围是 0 到 255用uint8存储最合适。转成张量之后如果要参与计算再转float32并归一化到 0 到 1。from PIL import Image import torchvision.transforms as T img Image.open(image.jpg) transform T.Compose([ T.ToTensor(), # 自动转成 float32 并归一化到 [0, 1] ]) tensor transform(img) print(tensor.dtype) # torch.float32ToTensor()这个转换做了三件事把 PIL 图像或 NumPy 数组转成张量把维度从 HWC 转成 CHW把数值从 0 到 255 归一化到 0 到 1。一步到位省去了手动转类型的麻烦。4.4 类型转换对计算图和梯度的影响这是一个容易被忽略的点。类型转换操作本身是可微的但有些转换会断开梯度。.to()和.float()这些方法如果输入需要梯度输出也需要梯度梯度会正常回传。但.detach()会断开梯度.long()转成整数类型之后梯度也没法回传了因为整数类型不支持梯度。x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) y x.float() # 梯度正常 y.sum().backward() print(x.grad) # tensor([1., 1., 1.]) x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) y x.long() # 转成整数梯度断裂 # y.sum().backward() # 报错element 0 of tensors does not require grad所以在模型内部不要随便把需要梯度的张量转成整数。如果只是临时用一下比如算索引用.detach().long()明确断开梯度避免意外。5. 常见报错与排查技巧实录5.1 典型报错速查表报错信息原因解决方法expected scalar type Float but found Double张量是 float64模型期望 float32x x.float()或x.to(torch.float32)expected scalar type Long but found Int索引是 int32框架期望 int64x x.long()或x.to(torch.int64)expected device cuda but got cpu张量在 CPU模型在 GPUx x.to(device)cant convert cuda:0 device type tensor to numpyGPU 张量直接转 NumPyx x.cpu().numpy()RuntimeError: result type Float cant be cast to the desired output type Long运算结果类型和期望不符检查运算双方的 dtype统一后再算ValueError: Tensor conversion requested dtype float32 for Tensor with dtype int64TensorFlow 中类型不匹配tf.cast(x, tf.float32)这张表里的报错我几乎每一个都遇到过。最常见的是第一个和第二个基本都是数据预处理没统一类型导致的。5.2 排查思路从报错信息定位到具体张量报错信息通常会告诉你期望什么类型、实际什么类型但不会告诉你哪个张量出了问题。在大模型里张量成千上万怎么定位我的做法是在报错的位置往前找看最近一次类型转换或者数据输入在哪里。如果是模型 forward 报错就在 forward 里逐层打印 dtype找到第一个不匹配的层。如果是数据加载报错就检查 Dataset 的__getitem__返回的类型。def forward(self, x): print(finput dtype: {x.dtype}) x self.layer1(x) print(fafter layer1: {x.dtype}) x self.layer2(x) print(fafter layer2: {x.dtype}) return x这种打印法虽然笨但很有效。找到出问题的层之后再回头看那一层的输入是从哪来的顺藤摸瓜就能找到根源。还有一个技巧是用torch.autograd.set_detect_anomaly(True)它会在反向传播出问题的时候给出更详细的堆栈信息。不过这个开关会拖慢训练速度只在调试的时候开。5.3 避坑经验数据管道中的类型统一策略踩了无数次坑之后我总结出一套数据管道的类型统一策略基本能避免 90% 的类型问题。第一步在 Dataset 的__getitem__里就把类型定死。图像转float32标签转int64不要等到 collate 或者模型里再转。class MyDataset(torch.utils.data.Dataset): def __getitem__(self, idx): image self.images[idx].astype(np.float32) # 明确转 float32 label np.int64(self.labels[idx]) # 明确转 int64 return torch.from_numpy(image), torch.tensor(label, dtypetorch.long)第二步在模型定义里第一层之前加一个类型检查或者转换。虽然多了一步但能兜底。class MyModel(nn.Module): def forward(self, x): x x.float() # 兜底转换 return self.net(x)第三步混合精度训练用amp不要手动转float16。amp会自动处理类型比手动靠谱。第四步定期用torch.typename()或者x.dtype检查关键张量的类型尤其是在数据增强、拼接、切片之后。这些操作有时候会改变类型比如torch.cat要求所有张量类型一致。注意torch.cat和torch.stack要求所有输入张量的 dtype 一致否则报错。如果从不同来源拿到的张量类型不同先统一再拼接。这个坑我在做多模态数据的时候踩过图像特征是float32文本特征是float64拼接的时候直接报错。6. 进阶话题量化与自动混合精度中的类型转换6.1 量化推理中的 int8 转换原理与实操量化是把float32的权重和激活值映射到int8的过程能大幅减少模型体积和推理时间。PyTorch 提供了动态量化和静态量化两种方式。动态量化最简单一行代码搞定import torch.quantization model MyModel() model.eval() quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear}, dtypetorch.qint8 )qint8是量化后的整数类型底层是int8。量化后的模型权重是int8计算的时候动态转成float32再算算完再转回去。这样既省了存储又保证了计算精度。静态量化更复杂需要校准数据但推理速度更快。它把激活值的范围也提前定好推理的时候直接用int8算。model.qconfig torch.quantization.get_default_qconfig(fbgemm) model_prepared torch.quantization.prepare(model) # 用校准数据跑几轮 for data in calib_loader: model_prepared(data) model_quantized torch.quantization.convert(model_prepared)量化里的类型转换核心是scale和zero_point两个参数。scale是浮点到整数的缩放因子zero_point是浮点 0 对应的整数值。转换公式是int_val round(float_val / scale) zero_point。反量化就是反过来。6.2 自动混合精度 AMP让框架帮你管类型torch.cuda.amp是 PyTorch 官方推荐的混合精度方案。它的核心是autocast上下文管理器和GradScaler。autocast会自动判断每个操作该用什么类型。矩阵乘法、卷积这些适合float16的操作自动转float16softmax、loss 计算这些需要精度的操作保持float32。你不需要手动转任何类型框架帮你搞定。GradScaler解决的是梯度下溢问题。float16能表示的最小正数约 6e-8很多梯度比这个还小直接变成 0。GradScaler先把 loss 放大比如乘以 65536算出来的梯度也放大更新之前再缩回去。这样梯度就不会下溢了。scaler GradScaler() for epoch in range(epochs): for data, target in dataloader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.unscale_(optimizer) # 可选裁剪梯度前先还原 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update()这套流程我用了两年多训练过各种规模的模型稳定性很好。唯一需要注意的是autocast区域里不要手动转float16会冲突。让框架自己管就行。6.3 bfloat16新一代训练类型的转换实践bfloat16是 Google 提出的格式指数位和float32一样是 8 位尾数只有 7 位。它的动态范围和float32一样不会溢出但精度比float16低。在 TPU 和 A100 之后的 GPU 上bfloat16的训练效果很好而且不需要GradScaler因为不会下溢。# 使用 bfloat16 的 autocast with torch.autocast(device_typecuda, dtypetorch.bfloat16): output model(data) loss criterion(output, target) loss.backward() optimizer.step()bfloat16的转换和float16类似用.to(torch.bfloat16)或者.bfloat16()。但要注意不是所有 GPU 都支持bfloat16需要 Ampere 架构之后的显卡。用之前先检查torch.cuda.is_bf16_supported()。我在 A100 上用bfloat16训练过一个 transformer相比float16加GradScaler的方案代码更简洁训练也更稳定loss 曲线平滑很多。如果你的硬件支持强烈推荐试试。7. 我个人的类型转换经验总结写了这么多最后分享几个我在实际项目里总结的小经验都是踩坑换来的。第一个经验在数据管道的出口统一类型不要在模型里到处转。模型里的类型转换会拖慢速度而且容易漏。在 Dataset 里把类型定死后面就省心了。第二个经验遇到类型报错先看报错信息里的期望类型和实际类型再往前找最近的数据来源。大部分类型问题都是数据预处理没做好而不是模型本身的问题。第三个经验混合精度训练用 amp不要手动转 float16。手动转容易溢出amp 自动处理边界情况省心得多。第四个经验int64 和 float32 是默认选择除非有明确理由换别的。int32 省的那点内存不值得带来的类型转换麻烦。float64 除非做科学计算否则没必要。第五个经验类型转换不是免费的频繁转换会拖慢训练。能在数据加载阶段转好的就不要留到模型里转。GPU 上的类型转换比 CPU 上快但也不是没有开销。第六个经验用torch.typename()或者x.dtype检查类型比猜靠谱。我见过太多人凭感觉写代码结果类型不对查半天。打印一下 dtype一秒钟的事。这些经验不一定适用于所有场景但至少能帮你少走一些弯路。张量类型转换这件事说难不难说简单也不简单。核心就是理解每种类型的用途和限制在合适的场景用合适的类型转换的时候注意精度和性能。多写多练慢慢就有感觉了。
返回列表