torch.linspace 是 PyTorch 中的一个函数,用于生成一个在指定区间内均匀分布的一维张量。这个函数在深度学习中常用于生成等间隔的数值序列,例如用于绘制图形、生成网格点等。

函数签名

torch.linspace(start, end, steps=100, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor

参数说明

  • start (float): 序列的起始值。
  • end (float): 序列的结束值。
  • steps (int, optional): 生成的样本数量。默认值为 100。
  • out (Tensor, optional): 输出张量。如果提供,结果将被写入这个张量中。
  • dtype (torch.dtype, optional): 返回张量的数据类型。如果未指定,则根据输入值推断数据类型。
  • layout (torch.layout, optional): 返回张量的内存布局。默认是 torch.strided
  • device (torch.device, optional): 返回张量所在的设备(如 CPU 或 GPU)。默认是当前设备。
  • requires_grad (bool, optional): 如果为 True,则返回的张量会记录操作以支持自动求导。默认是 False

返回值

  • 返回一个一维张量,包含从 startend 的等间隔数值序列,序列长度为 steps

示例

示例 1: 基本用法

import torch

# 生成从 0 到 10 的 5 个等间隔数值
tensor = torch.linspace(0, 10, steps=5)
print(tensor)

输出:

tensor([ 0.0000,  2.5000,  5.0000,  7.5000, 10.0000])

示例 2: 指定数据类型和设备

import torch

# 生成从 1 到 100 的 10 个等间隔数值,数据类型为 float32,设备为 GPU
tensor = torch.linspace(1, 100, steps=10, dtype=torch.float32, device='cuda')
print(tensor)

输出:

tensor([  1.0000,  12.0000,  23.0000,  34.0000,  45.0000,  56.0000,  67.0000,  78.0000,  89.0000, 100.0000], device='cuda:0')

示例 3: 使用 out 参数

import torch

# 创建一个空的输出张量
out_tensor = torch.empty(5)

# 将生成的序列写入 out_tensor
torch.linspace(0, 1, steps=5, out=out_tensor)
print(out_tensor)

输出:

tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])

注意事项

  • startend 可以是任意浮点数或整数。
  • steps 必须是一个正整数,表示生成的样本数量。
  • 如果 steps 为 1,则返回的张量将只包含 start 值。
  • dtype 参数可以控制返回张量的数据类型,例如 torch.float32torch.float64 等。

总结

torch.linspace 是一个非常有用的函数,特别适用于生成等间隔的数值序列。它在深度学习中的应用场景非常广泛,例如在生成网格点、绘制图形、初始化参数等方面都非常有用。