torch.tensor 用法
在 PyTorch 中,torch.tensor
是创建张量的核心函数。以下是详细用法指南:
一、基础用法
import torch
# 从 Python 列表/数组创建
data = [1, 2, 3]
tensor = torch.tensor(data) # 输出:tensor([1, 2, 3])
# 从 NumPy 数组创建(需先转 NumPy)
import numpy as np
numpy_arr = np.array([4, 5, 6])
tensor_from_np = torch.tensor(numpy_arr) # tensor([4, 5, 6])
二、关键参数详解
torch.tensor(
data, # 输入数据 (list/np.array 等)
dtype=None, # 指定数据类型 (torch.float32, torch.int64 等)
device=None, # 设备 ("cpu" 或 "cuda:0")
requires_grad=False, # 是否跟踪梯度
pin_memory=False # 是否固定内存(加速 GPU 传输)
)
三、常用场景示例
1. 指定数据类型
# 创建浮点型张量
float_tensor = torch.tensor([1, 2, 3], dtype=torch.float32) # tensor([1., 2., 3.])
# 创建布尔型张量
bool_tensor = torch.tensor([0, 1, 0], dtype=torch.bool) # tensor([False, True, False])
2. 设备选择 (CPU/GPU)
# 在 GPU 上创建张量 (需 CUDA 可用)
gpu_tensor = torch.tensor([7, 8, 9], device="cuda:0")
# 将 CPU 张量移动到 GPU
cpu_tensor = torch.tensor([10, 11, 12])
gpu_tensor = cpu_tensor.to("cuda:0")
3. 梯度跟踪
# 创建需要计算梯度的张量
x = torch.tensor([3.0], requires_grad=True)
y = x**2
y.backward() # 自动计算梯度
print(x.grad) # 输出:tensor([6.])
四、与其他创建方式的区别
方法 | 特点 | 示例 |
---|---|---|
torch.tensor() | 显式复制数据并推断类型 | torch.tensor([1, 2, 3]) |
torch.Tensor() | 类构造函数,默认 dtype=float32 | torch.Tensor([1, 2, 3]) → float |
torch.from_numpy() | 与 NumPy 共享内存 | torch.from_numpy(np_array) |
五、注意事项
- 数据复制行为:
torch.tensor()
会复制输入数据,若需要共享内存用torch.from_numpy()
- 隐式类型转换:
# 列表包含混合类型时会向上转型 mixed_data = [1, 2.0] # 会被转为 float
- 性能优化:对于大型数据集,优先使用
torch.utils.data.DataLoader
而非多次调用torch.tensor
六、高级技巧
# 创建未初始化张量(需后续填充数据)
empty_tensor = torch.tensor([], dtype=torch.float32)
empty_tensor.resize_(3, 4) # 调整形状为 3x4
# 从生成器创建张量
gen = (i*2 for i in range(5)) # 生成器表达式
tensor_from_gen = torch.tensor(list(gen)) # tensor([0, 2, 4, 6, 8])
如果需要创建特殊张量(全零/单位矩阵等),建议使用:
torch.zeros(2, 3) # 全零矩阵
torch.ones(3, 2) # 全一矩阵
torch.eye(4) # 单位矩阵
torch.randn(5, 5) # 标准正态分布随机数