PyTorch Tensor的创建、运算与GPU加速实战

发布时间:2026/7/5 12:00:45
PyTorch Tensor的创建、运算与GPU加速实战 1. PyTorch Tensor基础概念与创建方法Tensor是PyTorch中最核心的数据结构你可以把它理解为Numpy数组的升级版。想象一下Tensor就像是一个可以放在GPU上运行的超级数组它能帮我们快速完成各种数学运算。我第一次接触Tensor时发现它和Numpy的ndarray非常相似但多了一个超能力——GPU加速。创建Tensor的方法多种多样最基础的是使用torch.tensor()函数。比如你想把一个Python列表转换成Tensorimport torch data [[1, 2], [3, 4]] x torch.tensor(data) print(x)这里有个新手常踩的坑如果你给的列表维度不整齐比如[[1,2],[3]]PyTorch会直接报错。我刚开始就犯过这个错误调试了半天才发现是数据格式问题。PyTorch还提供了一系列便捷的初始化函数zeros torch.zeros(2, 3) # 全0矩阵 ones torch.ones(2, 3) # 全1矩阵 rand torch.rand(2, 3) # 0-1均匀分布随机数 randn torch.randn(2, 3) # 标准正态分布随机数在实际项目中我经常用torch.randn来初始化神经网络权重。记得设置随机种子保证结果可复现torch.manual_seed(42) # 设置随机种子 a torch.randn(2, 2) b torch.randn(2, 2) print(a b) # 每次运行结果相同2. Tensor的常用操作与运算掌握了Tensor的创建方法后我们来看看它能做什么。Tensor支持几乎所有你能想到的数学运算而且语法非常直观。2.1 基础数学运算加减乘除这些基础运算可以直接用运算符a torch.tensor([1, 2, 3]) b torch.tensor([4, 5, 6]) print(a b) # 逐元素相加 print(a * b) # 逐元素相乘矩阵乘法要用或torch.matmul()x torch.randn(2, 3) y torch.randn(3, 4) z x y # 矩阵乘法 print(z.shape) # 输出(2,4)2.2 广播机制PyTorch的广播机制和Numpy一样智能。比如你想把一个向量加到矩阵的每一行上matrix torch.ones(3, 4) vector torch.arange(4) result matrix vector # 自动广播 print(result)广播规则简单来说就是从最后一个维度开始比较要么维度大小相同要么其中一个为1或者其中一个维度不存在。我在实际项目中经常用广播来简化代码避免不必要的循环。2.3 形状操作改变Tensor形状是家常便饭常用的方法有x torch.arange(12) y x.view(3, 4) # 改变形状 z x.reshape(3, 4) # 功能类似view print(y.shape, z.shape)view和reshape的主要区别在于view要求数据在内存中是连续的而reshape会自动处理非连续情况。如果遇到view报错可以先用contiguous()方法。3. Tensor与Numpy的互操作PyTorch和Numpy可以无缝协作这在数据处理阶段特别有用。我经常先用Numpy处理原始数据再转成Tensor喂给模型。3.1 Tensor转Numpya torch.ones(3) b a.numpy() # Tensor转Numpy print(type(b)) # class numpy.ndarray需要注意的是如果Tensor在GPU上需要先移到CPUif torch.cuda.is_available(): a a.cpu().numpy()3.2 Numpy转Tensorimport numpy as np arr np.array([1, 2, 3]) tensor torch.from_numpy(arr) # Numpy转Tensor print(tensor)转换后的Tensor和原Numpy数组共享内存修改一个会影响另一个。这在某些情况下会导致难以发现的bug需要特别注意。4. GPU加速实战终于来到最激动人心的部分——GPU加速。PyTorch让GPU计算变得异常简单这也是它深受欢迎的重要原因。4.1 检查GPU可用性首先检查你的设备是否支持CUDAprint(torch.cuda.is_available()) # 输出True表示可用 print(torch.cuda.device_count()) # 查看GPU数量 print(torch.cuda.get_device_name(0)) # 查看GPU型号4.2 Tensor在CPU和GPU间移动把Tensor放到GPU上非常简单device torch.device(cuda if torch.cuda.is_available() else cpu) x torch.randn(3, 3).to(device) # 移动到GPU y torch.randn(3, 3).cuda() # 另一种写法从GPU移回CPUx_cpu x.cpu() # 移回CPU4.3 直接在GPU上创建Tensor为了获得最佳性能可以跳过CPU直接在GPU上创建Tensorgpu_tensor torch.zeros(3, 3, devicecuda)我在训练大型模型时发现直接在GPU上初始化参数能节省约10%的时间。对于超大规模训练这个优化非常值得。4.4 多GPU数据并行如果你的机器有多个GPU可以用DataParallel轻松实现数据并行model MyModel() if torch.cuda.device_count() 1: model torch.nn.DataParallel(model) model.to(device)这样模型会自动把数据切分到各个GPU上并行计算。我曾经用这个方法把训练速度提高了近3倍使用4块GPU。5. 性能优化技巧经过多次项目实践我总结了一些Tensor操作的性能优化经验尽量使用内置函数PyTorch的内置函数经过高度优化比自己用Python实现的要快得多。减少CPU-GPU数据传输频繁在CPU和GPU之间拷贝数据会严重影响性能。尽量保持数据在GPU上完成所有操作。使用原地操作带下划线的方法如add_()可以节省内存a torch.rand(3,3) a.add_(1) # 原地加1不创建新Tensor合理使用torch.no_grad()在不需要计算梯度的场景下使用with torch.no_grad(): # 这里面的操作不会跟踪梯度 y model(x)选择合适的精度大多数情况下float32足够用某些场景可以尝试float16来节省内存和计算量。记得第一次训练GAN模型时我因为没注意这些优化点训练速度比预期慢了近5倍。后来逐步应用这些技巧性能得到了显著提升。