1. 项目概述从“形状”到“内容”的桥梁在PyTorch的日常开发中我们经常遇到一个看似简单却至关重要的需求我需要一个和现有张量形状完全一样但内容全是1或全是0的新张量。比如在构建一个与输入尺寸相同的掩码mask时在初始化一个与梯度张量形状一致的累加器时或者在进行某些需要单位矩阵或零矩阵参与的运算时。手动计算尺寸并用torch.ones()或torch.zeros()创建当然可以但代码会显得冗长且容易出错尤其是在处理动态形状或复杂维度时。这时torch.ones_like和torch.zeros_like这两个函数就成了我们工具箱里的“瑞士军刀”它们能精准地复制一个张量的“形”并赋予其全新的“神”——全1或全0的值。这两个函数的核心价值在于“一致性”和“便捷性”。它们确保了新张量与参考张量在维度、大小上分毫不差同时省去了我们手动传递尺寸参数的麻烦。更重要的是它们能智能地继承参考张量的设备CPU或GPU、数据类型dtype等关键属性这在混合精度训练或多设备编程场景下尤为重要。一个错误的设备或数据类型匹配问题就可能导致程序崩溃或难以察觉的性能损失。因此深入理解并熟练运用这两个函数是写出健壮、高效PyTorch代码的基本功。无论你是刚入门的新手还是已经构建过复杂模型的老手重新审视这些基础工具背后的设计哲学和实用技巧总能带来新的收获。2. 函数深度解析不仅仅是“复制形状”2.1 函数签名与核心参数让我们先抛开简单的“创建全1/全0张量”的认知深入到函数的签名层面看看它们到底提供了哪些控制能力。torch.ones_like(input, *, dtypeNone, deviceNone, requires_gradFalse, layouttorch.strided) - Tensortorch.zeros_like(input, *, dtypeNone, deviceNone, requires_gradFalse, layouttorch.strided) - Tensor这里的input就是我们的参考张量。函数会返回一个与input形状相同的新张量。关键在于后面那些以*开头的关键字参数它们给了我们覆盖默认行为的权力dtype(数据类型)新张量的数据类型。默认行为是None这意味着新张量将继承input的dtype。这是最常用也最安全的方式能保证运算兼容。但有时我们需要显式指定例如将float32的参考张量对应的掩码创建为bool类型以节省内存mask torch.ones_like(img_tensor, dtypetorch.bool)。device(设备)新张量所在的设备CPU或特定的GPU。默认None同样表示继承自input。这在数据并行或模型分段部署时非常有用可以确保中间张量被创建在正确的设备上避免昂贵的设备间数据传输。requires_grad(是否需要梯度)一个布尔值决定新张量是否参与自动微分计算。默认是False。这一点极其重要如果你创建的这个张量比如一个可学习的偏置项初始化需要被优化器更新就必须将其设为Truebias torch.zeros_like(some_tensor, requires_gradTrue)。反之如果只是一个临时掩码或常量保持False可以减少计算图的内存开销。layout(内存布局)通常我们使用默认的torch.strided跨步布局即可它适用于绝大多数稠密张量操作。在处理稀疏张量等特殊场景时才会需要调整。注意*在函数签名中表示其后的参数必须通过关键字参数的形式指定不能通过位置参数传递。这提高了代码的可读性避免了因参数顺序导致的错误。2.2 与手动创建方式的对比为了更直观地理解*_like系列函数的优势我们将其与手动创建进行对比。假设我们有一个从数据加载器中获取的批处理张量batch其形状为[32, 3, 224, 224]批次大小323通道224x224图像设备在GPU上cuda:0数据类型为torch.float32。目标创建一个与之形状相同的全1张量作为权重掩码。方法一手动创建繁琐且易错mask torch.ones(batch.size(0), batch.size(1), batch.size(2), batch.size(3), dtypebatch.dtype, devicebatch.device) # 或者 mask torch.ones(batch.shape, dtypebatch.dtype, devicebatch.device)你需要显式地获取batch的每一个维度或者使用.shape属性然后还要记得传递dtype和device。代码冗长且在模型输入维度动态变化时如可变序列长度维护起来很麻烦。方法二使用ones_like简洁且鲁棒mask torch.ones_like(batch)一行代码搞定。无论batch的形状如何变化比如变成[16, 1, 512, 512]无论它是在CPU还是GPU上无论它是float32还是float16mask都会自动保持同步。这种写法将“形状和属性一致性”这个责任完全交给了框架大大减少了人为错误的可能性也让代码意图更加清晰——“我要一个和batch一样的全1张量”。2.3 继承机制的内部逻辑理解默认的继承机制能帮助我们避免陷阱。当你不指定dtype,device等参数时函数内部大致执行了如下逻辑def ones_like_impl(input, **kwargs): # 获取或使用用户提供的参数若未提供则使用input的属性 target_dtype kwargs.get(dtype, input.dtype) target_device kwargs.get(device, input.device) target_requires_grad kwargs.get(requires_grad, False) # ... 其他属性 # 调用底层的Tensor构造函数创建指定形状和属性的全1张量 return torch.ones(input.size(), dtypetarget_dtype, devicetarget_device, requires_gradtarget_requires_grad)这解释了为什么它如此方便。但在复杂场景下我们有时需要打破这种继承。例如在混合精度训练中模型参数可能是float16半精度但某些中间计算如损失值需要保持float32单精度以获得足够的数值精度。这时创建一个与float16张量形状相同但类型为float32的零张量来累加损失就需要显式指定dtype。3. 核心应用场景与实战技巧3.1 掩码Mask生成这是ones_like和zeros_like最经典的应用。在自然语言处理NLP和计算机视觉CV中掩码用于标识有效数据区域或需要特别处理的部分。场景一序列填充掩码NLP在Transformer等模型中需要对变长序列进行批处理短序列会被填充pad到同一长度。我们需要一个掩码来告诉模型哪些位置是真实的token哪些是填充的。# 假设 sequences 是经过填充后的张量形状为 [batch_size, seq_len] # pad_token_id 是填充符的ID比如0 padding_mask (sequences ! pad_token_id) # 得到一个布尔掩码 # 但在某些注意力计算中需要将填充位置的值设为极小的负数如-1e9以便softmax后权重为0 # 我们可以先创建一个全1的矩阵然后将填充位置设为0 attention_mask torch.ones_like(sequences, dtypetorch.float32) # 继承形状和设备 attention_mask attention_mask.masked_fill(padding_mask 0, 0.0) # 填充位置置0 # 或者在注意力分数计算前直接创建一个与分数矩阵形状相同的掩码矩阵进行加法 scores some_attention_function(q, k) # scores形状: [batch, heads, seq_len, seq_len] # 创建一个与scores形状相同的掩码并将需要屏蔽的位置设为很大的负数 extended_padding_mask padding_mask.unsqueeze(1).unsqueeze(2) # 扩展维度以匹配scores negative_mask torch.zeros_like(scores) negative_mask negative_mask.masked_fill(extended_padding_mask 0, -1e9) scores scores negative_mask这里ones_like和zeros_like确保了掩码与目标张量形状的绝对一致避免了维度不匹配的错误。场景二图像区域掩码CV在图像分割或目标检测中可能需要根据边界框生成一个二值掩码。def create_box_mask(image_tensor, box_coords): image_tensor: 输入图像张量 [C, H, W] box_coords: 边界框坐标 [x_min, y_min, x_max, y_max] # 首先创建一个和图像空间维度H, W相同的全零掩码 # 注意这里参考的是图像的空间维度而不是整个image_tensor h, w image_tensor.shape[-2:] mask torch.zeros((h, w), dtypetorch.bool, deviceimage_tensor.device) # 显式指定bool类型节省内存 # 或者如果你想创建一个与image_tensor通道数也一致的掩码较少见 # mask torch.zeros_like(image_tensor, dtypetorch.bool) # 形状为[C, H, W] x_min, y_min, x_max, y_max box_coords mask[y_min:y_max, x_min:x_max] True # 将框内区域置为True return mask这个例子展示了如何灵活选择参考张量。我们并不总是需要复制整个输入张量的所有属性有时只需要其空间维度。此时直接使用torch.zeros((h, w), ...)可能比zeros_like更清晰。但如果你已经有一个表示单通道掩码的张量template那么zeros_like(template)就是最直接的选择。3.2 初始化与累加器在模型训练或自定义算子中经常需要初始化与某个中间变量形状相同的张量。场景一自定义损失函数中的累加假设你实现了一个需要逐元素计算权重的损失函数。def weighted_mse_loss(pred, target, weight_map): # weight_map 形状与 pred、target 相同 diff pred - target squared_diff diff ** 2 # 我们需要一个与squared_diff形状相同的临时张量来存放加权后的结果 # 使用 zeros_like 进行初始化是最安全的方式 weighted_squared_diff torch.zeros_like(squared_diff) weighted_squared_diff squared_diff * weight_map loss weighted_squared_diff.mean() return loss虽然这里可以直接写squared_diff * weight_map但使用zeros_like初始化weighted_squared_diff在逻辑上更清晰尤其是在后续可能需要进行更复杂的逐元素操作时能确保有一个干净的起点。场景二梯度累加或动量计算在实现优化算法如SGD with Momentum时需要维护一个与参数相同形状的“速度”张量。class SimpleSGDWithMomentum: def __init__(self, params, lr0.01, momentum0.9): self.params list(params) self.lr lr self.momentum momentum # 关键步骤为每个参数初始化一个速度缓冲区全零 self.velocity [torch.zeros_like(p) for p in self.params] def step(self): for p, v in zip(self.params, self.velocity): if p.grad is None: continue # 更新速度: v momentum * v - lr * grad v.mul_(self.momentum).add_(p.grad, alpha-self.lr) # 更新参数: p p v p.add_(v)这里torch.zeros_like(p)确保了每个速度项v与其对应的参数p在形状、设备、数据类型上完全一致这是优化器正确工作的前提。如果手动创建很容易遗漏device或dtype导致运行时错误。3.3 广播机制下的协同工作PyTorch的广播机制非常强大但有时为了代码的明确性或性能我们需要手动创建具有特定形状的张量。*_like函数在这里也能发挥作用。场景确保明确的广播形状假设我们有一个特征张量features形状为[B, C, H, W]和一个每通道的缩放因子scale形状为[C]。我们想对每个通道进行缩放。# 方法A依赖自动广播简洁 scaled_features features * scale.view(1, -1, 1, 1) # 将scale变形为[1, C, 1, 1] # 方法B使用ones_like创建显式的缩放矩阵有时更清晰 # 创建一个与features形状相同的全1张量 scale_matrix torch.ones_like(features) # 将每一通道的缩放因子填充进去。这里需要一些维度操作。 for c in range(features.size(1)): scale_matrix[:, c, :, :] scale[c] scaled_features features * scale_matrix方法B显然更冗长且效率可能更低但它将“缩放”这个操作可视化成了一个与原始数据形状一致的矩阵在某些教学或调试场景下更有助于理解。而ones_like为创建这个“模板矩阵”提供了极大便利。在实际高性能代码中我们当然首选方法A但理解方法B背后的思维过程很重要。4. 高级用法、性能考量与常见陷阱4.1 内存布局Layout与稀疏张量绝大多数情况下我们处理的是稠密张量torch.strided布局。ones_like和zeros_like默认创建的就是这种布局。但PyTorch也支持稀疏张量如torch.sparse_coo。如果你有一个稀疏张量作为输入*_like函数的行为会稍有不同。import torch # 创建一个稀疏张量 i torch.tensor([[0, 1, 1], [2, 0, 2]]) # 索引 v torch.tensor([3, 4, 5], dtypetorch.float32) # 值 sparse_input torch.sparse_coo_tensor(i, v, (3, 3)) # 尝试创建 like 张量 dense_like torch.ones_like(sparse_input) # 这会发生什么 print(dense_like) # 输出一个形状为(3,3)的稠密全1张量。 # tensor([[1., 1., 1.], # [1., 1., 1.], # [1., 1., 1.]]) # 如果你想创建一个具有相同稀疏结构的全1张量需要指定layout # 但 ones_like 的 layout 参数主要用于指定输出布局不能直接复制输入的非稠密布局结构。 # 对于稀疏张量更常见的操作是创建一个具有相同索引但值全为1的新稀疏张量。 new_values torch.ones_like(v) # 复制值的形状和属性 sparse_ones torch.sparse_coo_tensor(i, new_values, sparse_input.size())关键点在于*_like函数复制的是形状size和属性dtype, device等而不是稀疏结构。对于稀疏张量它返回的是一个具有相同形状的稠密张量。如果你需要操作稀疏结构本身需要直接处理稀疏张量的indices()和values()属性。4.2 原地操作In-place与梯度计算这是一个容易踩坑的地方。torch.ones_like和torch.zeros_like总是返回一个新的张量。它们不是原地操作。这通常是我们所期望的。但当你需要修改一个现有张量的值使其变为全1或全0时应该使用.fill_()或.zero_()方法。x torch.tensor([[1., 2.], [3., 4.]], requires_gradTrue) y x * 2 # 错误做法这创建了一个新的张量打断了计算图 mask torch.ones_like(y) result y mask # mask是一个新的叶子节点与x的计算图无关 # 正确做法如果需要新张量上述做法在需要独立掩码时是对的。 # 如果需要原地重置y谨慎这会破坏梯度 # y.zero_() # 将y的所有元素置0并且这个操作会被记录在计算图中但通常不是我们想要的。 # 更常见的需求创建一个与y形状相同、需要梯度的全1参数 learnable_scale torch.ones_like(y, requires_gradTrue) # 这是一个新的可学习参数 result y * learnable_scale # learnable_scale的梯度可以被计算关于requires_grad的决策至关重要requires_gradFalse(默认)适用于掩码、常量、临时缓冲区。不参与梯度计算节省内存。requires_gradTrue适用于你希望其值通过梯度下降进行优化的张量例如自定义层的初始化参数。4.3 设备Device一致性错误排查在多GPU训练或混合设备环境中device不一致是常见的运行时错误来源。*_like函数的自动继承特性是避免此类问题的第一道防线。典型错误场景# 假设模型在GPU 0上 model model.to(cuda:0) input_data input_data.to(cuda:0) # 在某个函数内部不小心创建了一个在CPU上的掩码 def some_function(tensor): # 错误没有指定device默认在CPU上创建 bad_mask torch.ones_like(tensor.cpu()) # 或者更隐蔽的tensor之前被移到过CPU # 后续运算会崩溃 # result tensor * bad_mask # RuntimeError: Expected all tensors to be on the same device... # 正确依赖输入tensor的device good_mask torch.ones_like(tensor) # 自动继承 cuda:0 return tensor * good_mask排查技巧当遇到设备不匹配错误时首先检查错误信息中提到的张量然后使用.device属性查看每个相关张量所在的设备。确保所有参与运算的张量都在同一设备上。利用*_like函数是保持设备一致性的最佳实践。4.4 性能微考量对于大规模张量torch.ones_like和torch.zeros_like的调用开销极低它们本质上是对底层内存分配和初始化操作的封装。但在极端性能敏感的循环中仍需注意避免在循环内重复创建如果形状不变的张量在循环中被反复创建应该将其提到循环外部。# 低效 for batch in dataloader: mask torch.ones_like(batch) # 每次循环都新建 # ... use mask # 高效 (假设所有batch形状相同) sample_batch next(iter(dataloader)) mask_template torch.ones_like(sample_batch) for batch in dataloader: # 直接使用预先创建好的mask_template或如果需要复用确保其device/dtype一致 # 注意如果batch形状变化此方法不适用。 result batch * mask_template使用torch.empty_likefill_在极其罕见的、对初始化性能有苛刻要求的场景下可以先分配未初始化内存再填充。但ones_like和zeros_like内部已经做了高度优化绝大多数情况下无需手动这样做且empty_like后不立即填充会导致值未定义风险很高。# 一般不推荐除非有非常确切的性能瓶颈证据 a torch.empty_like(reference_tensor) a.fill_(1) # 或 a.zero_()5. 与其他相似函数的对比与选择PyTorch提供了多个用于创建张量的函数理解它们的区别有助于做出正确选择。函数核心功能与*_like的关键区别适用场景torch.ones()/torch.zeros()根据给定的**尺寸size**创建全1/全0张量。需要显式指定尺寸元组如torch.ones((2,3))。不自动继承其他属性。已知固定尺寸或需要创建与现有张量无关的新张量。torch.ones_like()/torch.zeros_like()根据给定的输入张量创建全1/全0张量。自动继承输入张量的形状并可选择性地继承或覆盖dtype,device等属性。强依赖现有张量形状且希望保持属性一致最常用。torch.full()/torch.full_like()创建填充指定标量值的张量。full(size, fill_value)full_like(input, fill_value)。比ones/zeros更通用。需要创建填充特定常数的张量例如全为5或全为-inf的矩阵。torch.empty()/torch.empty_like()分配未初始化内存的张量。内容为任意值取决于内存状态速度最快但不安全。性能瓶颈确证且会立即覆盖所有数据的情况。torch.rand()/torch.rand_like()创建服从[0,1)均匀分布的随机张量。用于初始化随机权重。模型权重初始化。torch.randn()/torch.randn_like()创建服从标准正态分布N(0,1)的随机张量。用于初始化符合正态分布的权重。模型权重初始化更常用。选择指南当你有一个参考张量并且新张量需要与其形状一致时首选*_like系列函数。这是最安全、最简洁的方式。如果形状是已知的简单字面量如(256, 256)使用torch.ones()或torch.zeros()更直接。如果需要填充非0非1的值使用torch.full_like()。除非在精心优化的代码段中且能保证安全否则避免使用torch.empty_like()。6. 常见问题与调试实录在实际使用中即使是这样基础的函数也可能遇到意想不到的问题。下面记录了几个典型案例和排查思路。问题1RuntimeError: expected device cuda:0 but got device cpu现象在GPU训练时突然报错设备不匹配。排查检查报错行代码找到涉及的所有张量。在每个张量后添加.device打印或者使用调试器查看。很可能某个torch.ones()或torch.zeros()创建时没有指定device导致默认创建在CPU上。而与之运算的另一个张量在GPU上。解决将出问题的创建语句改为torch.ones_like(gpu_tensor)或者显式指定devicegpu_tensor.device。问题2创建的全1张量参与计算后梯度为None现象自定义了一个损失函数里面用torch.ones_like创建了一个权重矩阵但训练时发现这个权重矩阵的梯度一直是None无法被优化器更新。排查检查创建语句weight torch.ones_like(input_tensor)。默认情况下requires_gradFalse。即使input_tensor需要梯度ones_like创建的新张量也不需要。解决如果这个权重是需要学习的参数必须显式设置requires_gradTrueweight torch.ones_like(input_tensor, requires_gradTrue)。问题3使用*_like后内存占用异常增加现象在循环中内存使用量不断上升。排查检查循环中是否持续用*_like创建新的大张量而旧的张量由于被变量引用而无法释放。例如buffer torch.zeros_like(large_tensor)在循环内每次迭代都创建一个新的buffer并赋值给同一个变量名。如果旧的张量还被其他对象引用比如添加到列表中它就不会被垃圾回收。解决如果可能复用张量在循环外创建一次在循环内使用.zero_()或.fill_()重置内容。确保不再需要的张量的引用被解除以便Python垃圾回收器或PyTorch的CUDA缓存管理器可以释放内存。对于非常大的张量考虑使用with torch.no_grad():上下文管理器来创建避免不必要的梯度计算图构建。问题4期望布尔掩码却得到了浮点型张量现象进行逻辑运算或作为索引时类型不匹配。排查mask torch.ones_like(float_tensor)创建的是与float_tensor同类型的浮点张量其值为1.0而不是布尔值True。直接将其用作布尔索引如tensor[mask]会报错或产生非预期结果。解决明确指定dtypetorch.bool。# 创建一个与float_tensor形状相同的全True布尔掩码 bool_mask torch.ones_like(float_tensor, dtypetorch.bool) # 或者更直接地从比较操作中获得布尔掩码 bool_mask (float_tensor 0.5)掌握torch.ones_like和torch.zeros_like远不止是记住两个API。它关乎你如何以一种一致、可靠且高效的方式思考张量的创建与协同工作。在PyTorch的动态图世界里让这些基础工具成为你肌肉记忆的一部分能让你更专注于算法逻辑本身而非琐碎的维度匹配或设备管理问题。下次当你需要一个新的张量时先问问自己“有没有一个现成的张量可以作为形状和属性的模板”如果有那么*_like函数很可能就是最优雅的解决方案。