ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

PyTorch中ones_like与zeros_like函数:高效创建形状匹配张量的核心技术

2026/8/13 10:19:11 拓冰建站 浏览量
PyTorch中ones_like与zeros_like函数:高效创建形状匹配张量的核心技术

1. 项目概述:从“形状”到“内容”的桥梁

在PyTorch的日常开发中,我们经常遇到一个看似简单却至关重要的需求:我需要一个和现有张量形状完全一样,但内容全是1或全是0的新张量。比如,在构建一个与输入尺寸相同的掩码(mask)时,在初始化一个与梯度张量形状一致的累加器时,或者在进行某些需要单位矩阵或零矩阵参与的运算时。手动计算尺寸并用torch.ones()torch.zeros()创建当然可以,但代码会显得冗长且容易出错,尤其是在处理动态形状或复杂维度时。这时,torch.ones_liketorch.zeros_like这两个函数就成了我们工具箱里的“瑞士军刀”,它们能精准地复制一个张量的“形”,并赋予其全新的“神”——全1或全0的值。

这两个函数的核心价值在于“一致性”和“便捷性”。它们确保了新张量与参考张量在维度、大小上分毫不差,同时省去了我们手动传递尺寸参数的麻烦。更重要的是,它们能智能地继承参考张量的设备(CPU或GPU)、数据类型(dtype)等关键属性,这在混合精度训练或多设备编程场景下尤为重要。一个错误的设备或数据类型匹配问题,就可能导致程序崩溃或难以察觉的性能损失。因此,深入理解并熟练运用这两个函数,是写出健壮、高效PyTorch代码的基本功。无论你是刚入门的新手,还是已经构建过复杂模型的老手,重新审视这些基础工具背后的设计哲学和实用技巧,总能带来新的收获。

2. 函数深度解析:不仅仅是“复制形状”

2.1 函数签名与核心参数

让我们先抛开简单的“创建全1/全0张量”的认知,深入到函数的签名层面,看看它们到底提供了哪些控制能力。

torch.ones_like(input, *, dtype=None, device=None, requires_grad=False, layout=torch.strided) -> Tensortorch.zeros_like(input, *, dtype=None, device=None, requires_grad=False, layout=torch.strided) -> Tensor

这里的input就是我们的参考张量。函数会返回一个与input形状相同的新张量。关键在于后面那些以*开头的关键字参数,它们给了我们覆盖默认行为的权力:

  1. dtype(数据类型):新张量的数据类型。默认行为是None,这意味着新张量将继承inputdtype。这是最常用也最安全的方式,能保证运算兼容。但有时我们需要显式指定,例如将float32的参考张量对应的掩码创建为bool类型以节省内存:mask = torch.ones_like(img_tensor, dtype=torch.bool)
  2. device(设备):新张量所在的设备(CPU或特定的GPU)。默认None同样表示继承自input。这在数据并行或模型分段部署时非常有用,可以确保中间张量被创建在正确的设备上,避免昂贵的设备间数据传输。
  3. requires_grad(是否需要梯度):一个布尔值,决定新张量是否参与自动微分计算。默认是False。这一点极其重要!如果你创建的这个张量(比如一个可学习的偏置项初始化)需要被优化器更新,就必须将其设为Truebias = torch.zeros_like(some_tensor, requires_grad=True)。反之,如果只是一个临时掩码或常量,保持False可以减少计算图的内存开销。
  4. layout(内存布局):通常我们使用默认的torch.strided(跨步布局)即可,它适用于绝大多数稠密张量操作。在处理稀疏张量等特殊场景时才会需要调整。

注意*在函数签名中表示其后的参数必须通过关键字参数的形式指定,不能通过位置参数传递。这提高了代码的可读性,避免了因参数顺序导致的错误。

2.2 与手动创建方式的对比

为了更直观地理解*_like系列函数的优势,我们将其与手动创建进行对比。

假设我们有一个从数据加载器中获取的批处理张量batch,其形状为[32, 3, 224, 224](批次大小32,3通道,224x224图像),设备在GPU上(cuda:0),数据类型为torch.float32

目标:创建一个与之形状相同的全1张量作为权重掩码。

方法一:手动创建(繁琐且易错)

mask = torch.ones(batch.size(0), batch.size(1), batch.size(2), batch.size(3), dtype=batch.dtype, device=batch.device) # 或者 mask = torch.ones(batch.shape, dtype=batch.dtype, device=batch.device)

你需要显式地获取batch的每一个维度,或者使用.shape属性,然后还要记得传递dtypedevice。代码冗长,且在模型输入维度动态变化时(如可变序列长度),维护起来很麻烦。

方法二:使用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(), dtype=target_dtype, device=target_device, requires_grad=target_requires_grad)

这解释了为什么它如此方便。但在复杂场景下,我们有时需要打破这种继承。例如,在混合精度训练中,模型参数可能是float16(半精度),但某些中间计算(如损失值)需要保持float32(单精度)以获得足够的数值精度。这时,创建一个与float16张量形状相同但类型为float32的零张量来累加损失,就需要显式指定dtype

3. 核心应用场景与实战技巧

3.1 掩码(Mask)生成

这是ones_likezeros_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, dtype=torch.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_likezeros_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), dtype=torch.bool, device=image_tensor.device) # 显式指定bool类型节省内存 # 或者,如果你想创建一个与image_tensor通道数也一致的掩码(较少见) # mask = torch.zeros_like(image_tensor, dtype=torch.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, lr=0.01, momentum=0.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在形状、设备、数据类型上完全一致,这是优化器正确工作的前提。如果手动创建,很容易遗漏devicedtype,导致运行时错误。

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_likezeros_like默认创建的就是这种布局。但PyTorch也支持稀疏张量(如torch.sparse_coo)。如果你有一个稀疏张量作为输入,*_like函数的行为会稍有不同。

import torch # 创建一个稀疏张量 i = torch.tensor([[0, 1, 1], [2, 0, 2]]) # 索引 v = torch.tensor([3, 4, 5], dtype=torch.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_liketorch.zeros_like总是返回一个新的张量。它们不是原地操作。这通常是我们所期望的。但当你需要修改一个现有张量的值,使其变为全1或全0时,应该使用.fill_().zero_()方法。

x = torch.tensor([[1., 2.], [3., 4.]], requires_grad=True) 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_grad=True) # 这是一个新的可学习参数 result = y * learnable_scale # learnable_scale的梯度可以被计算

关于requires_grad的决策至关重要:

  • requires_grad=False(默认):适用于掩码、常量、临时缓冲区。不参与梯度计算,节省内存。
  • requires_grad=True:适用于你希望其值通过梯度下降进行优化的张量,例如自定义层的初始化参数。

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_liketorch.zeros_like的调用开销极低,它们本质上是对底层内存分配和初始化操作的封装。但在极端性能敏感的循环中,仍需注意:

  1. 避免在循环内重复创建:如果形状不变的张量在循环中被反复创建,应该将其提到循环外部。

    # 低效 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
  2. 使用torch.empty_like+fill_:在极其罕见的、对初始化性能有苛刻要求的场景下,可以先分配未初始化内存,再填充。但ones_likezeros_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. 常见问题与调试实录

在实际使用中,即使是这样基础的函数,也可能遇到意想不到的问题。下面记录了几个典型案例和排查思路。

问题1:RuntimeError: expected device cuda:0 but got device cpu

  • 现象:在GPU训练时,突然报错设备不匹配。
  • 排查
    1. 检查报错行代码,找到涉及的所有张量。
    2. 在每个张量后添加.device打印,或者使用调试器查看。
    3. 很可能某个torch.ones()torch.zeros()创建时没有指定device,导致默认创建在CPU上。而与之运算的另一个张量在GPU上。
  • 解决:将出问题的创建语句改为torch.ones_like(gpu_tensor),或者显式指定device=gpu_tensor.device

问题2:创建的全1张量参与计算后,梯度为None

  • 现象:自定义了一个损失函数,里面用torch.ones_like创建了一个权重矩阵,但训练时发现这个权重矩阵的梯度一直是None,无法被优化器更新。
  • 排查
    1. 检查创建语句:weight = torch.ones_like(input_tensor)
    2. 默认情况下requires_grad=False。即使input_tensor需要梯度,ones_like创建的新张量也不需要。
  • 解决:如果这个权重是需要学习的参数,必须显式设置requires_grad=Trueweight = torch.ones_like(input_tensor, requires_grad=True)

问题3:使用*_like后内存占用异常增加

  • 现象:在循环中,内存使用量不断上升。
  • 排查
    1. 检查循环中是否持续用*_like创建新的大张量,而旧的张量由于被变量引用而无法释放。
    2. 例如: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])会报错或产生非预期结果。
  • 解决:明确指定dtype=torch.bool
    # 创建一个与float_tensor形状相同的全True布尔掩码 bool_mask = torch.ones_like(float_tensor, dtype=torch.bool) # 或者,更直接地,从比较操作中获得布尔掩码 bool_mask = (float_tensor > 0.5)

掌握torch.ones_liketorch.zeros_like,远不止是记住两个API。它关乎你如何以一种一致、可靠且高效的方式思考张量的创建与协同工作。在PyTorch的动态图世界里,让这些基础工具成为你肌肉记忆的一部分,能让你更专注于算法逻辑本身,而非琐碎的维度匹配或设备管理问题。下次当你需要一个新的张量时,先问问自己:“有没有一个现成的张量可以作为形状和属性的模板?”如果有,那么*_like函数很可能就是最优雅的解决方案。