PyTorch 基本使用创建、计算、转换、拼接、索引、形状、运算、加载保存等万字详解

请一直在路上

已于 2024-08-04 14:30:26 修改

阅读量1.7k

点赞数 1

CC 4.0 BY-SA版权

分类专栏：机器学习文章标签： numpy python 数据分析

于 2022-09-21 14:00:26 首次发布

本文链接：https://blog.youkuaiyun.com/qq_35240689/article/details/126967606

机器学习专栏收录该内容

9 篇文章

订阅专栏

本文详细介绍了 PyTorch 中张量的基本操作，包括创建、计算、类型转换、拼接、索引、形状调整等，并提供了模型保存和加载的方法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一、张量基本创建方式。
二、张量数值计算。
三、张量类型转换
四、张量拼接操作
五、张量的索引操作
六、张量的形状操作
七、张量运算函数
八、模型保存加载

一、张量基本创建方式。

1. 张量的基本概念

1. PyTorch 是一个 Python 深度学习框架，它将数据封装成张量（Tensor）来进行运算。
2. PyTorch 中的张量就是元素为同一种数据类型的多维矩阵。
3. PyTorch 中，张量以 "类" 的形式封装起来，对张量的一些运算、处理的方法被封装在类中。

2. 张量的基本创建

torch.tensor 根据指定数据创建张量

# 1. 根据已有的数据创建张量
data = torch.tensor(10)

# 2 使用numpy数组来创建张量
data = np.random.randn(2, 3)
data = torch.tensor(data)

# 3 使用list列表创建张量
data = [[10., 20., 30.], [40., 50., 60.]]
data = torch.tensor(data)

torch.Tensor 根据形状创建张量, 其也可用来创建指定数据的张量

# 2.1 创建2行3列的张量
data = torch.Tensor(2, 3)

# 2.2 可以创建指定值的张量、注意: 传递列表[]
data = torch.Tensor([2, 3])
data = torch.Tensor([10])

torch.IntTensor、torch.FloatTensor、torch.DoubleTensor 创建指定数据类型的张量

# 前面创建的张量都是使用默认类型或者元素类型
# 创建一个 int32 类型的张量
data = torch.IntTensor(2, 3)

# torch.ShortTensor(2, 3)  # 表示创建的是 int16 张量
# torch.LongTensor(2, 3)  # 表示创建的是 int32 张量
# torch.FloatTensor(2, 3)  # 表示创建的是 float32 张量

# 注意: 如果创建指定类型的张量，但是传递的数据不匹配，会发生类型转换
data = torch.IntTensor([2.5, 3.5])

3. 创建线性和随机张量

torch.arange 和 torch.linspace 创建线性张量

# 1.1 创建指定步长的张量
# (开始值，结束值，步长)
data = torch.arange(0, 10, 2)

# 1.2 在指定区间指定元素个数
# (开始值，结束值，创建元素的个数)
data = torch.linspace(0, 11, 10)

torch.randn 创建随机张量，torch.random.init_seed 和 torch.random.manual_seed 随机种子设置

# 2.1 创建随机张量
data = torch.randn(2, 3)

print('随机数种子:', torch.random.initial_seed())
# 随机种子设置
torch.random.manual_seed(100)
print('随机数种子:', torch.random.initial_seed())

4. 创建01张量和指定值张量

torch.ones 和 torch.ones_like 创建全1张量

# 1. 创建指定形状全0张量
data = torch.zeros(2, 3)

# 2. 根据张量形状创建全0张量
data = torch.zeros_like(data)

torch.zeros 和 torch.zeros_like 创建全0张量

# 1. 创建指定形状全0张量
data = torch.ones(2, 3)

# 2. 根据张量形状创建全0张量
data = torch.ones_like(data)

torch.full 和 torch.full_like 创建全为指定值张量

# 1. 创建指定形状指定值的张量
data = torch.full([2, 3], 10)

# 2. 根据张量形状创建指定值的张量
data = torch.full_like(data, 20)

4. 张量元素类型转换

tensor.type(torch.DoubleTensor)

data = torch.full([2, 3], 10)
print(data.dtype)

# 1. 将 data 元素类型转换为 float64 类型，第一种方法
data = data.type(torch.DoubleTensor)
print(data.dtype)

torch.double()

# 2. 第二种方法
data = data.double()
print(data.dtype)

二、张量数值计算。

1. 掌握张量基本运算

基本运算中，包括 add、sub、mul、div、neg 等函数, 以及 add_、sub_、mul_、div_、neg_，其中带下划线的版本为修改原数据。

import numpy as np
import torch

data = torch.randint(0, 10, [2, 3])
print(data)

# 1. 不修改原数据
new_data = data.add(10)  # 等价 new_data = data + 10
print(new_data)

# 2. 直接修改原数据
data.add_(10)  # 等价 data += 10
print(data)

# 3. 其他函数
print(data.sub(100))
print(data.mul(100))
print(data.div(100))
print(data.neg())

2. 阿达玛积

阿达玛积指的是矩阵对应位置的元素相乘.

import torch

data1 = torch.tensor([[1, 2], [3, 4]])
data2 = torch.tensor([[5, 6], [7, 8]])

# 第一种方式 mul
data = torch.mul(data1, data2)
print(data)

# 第二种方式 *
data = data1 * data2
print(data)

3. 点积运算

点积运算要求第一个矩阵 shape: (n, m)，第二个矩阵 shape: (m, p), 两个矩阵点积运算 shape 为: (n, p)。

import torch

data1 = torch.tensor([[1, 2], [3, 4], [5, 6]])
data2 = torch.tensor([[5, 6], [7, 8]])

# 第一种方式 @
data = data1 @ data2
print(data)
print('-' * 50)

# 第二种方式 mm
data = torch.mm(data1, data2)
print(data)
print('-' * 50)

# 第三种方式 matmul
data = torch.matmul(data1, data2)
print(data)
print('-' * 50)

torch.mm 用于进行两个矩阵点乘运算, 要求输入的矩阵为2维

data1 = torch.tensor([[1, 2], [3, 4], [5, 6]])
data2 = torch.tensor([[5, 6], [7, 8]])

data = torch.mm(data1, data2)
print(data)

torch.bmm 用于批量进行矩阵点乘运算, 要求输入的矩阵为3维

data1 = torch.randn(3, 4, 5)
data2 = torch.randn(3, 5, 8)

data = torch.bmm(data1, data2)
print(data.shape)

torch.matmul 用于进行两个矩阵点乘运算, 不要求输入的矩阵为2维，对数输入的 shape 不同的张量, 对应的最后几个维度必须符合矩阵运算规则

print(torch.matmul(torch.randn(3, 4, 5), torch.randn(5, 4)).shape)
print(torch.matmul(torch.randn(5, 4), torch.randn(3, 4, 5)).shape)

4. 指定运算设备

PyTorch 默认会将张量创建在 CPU 控制的内存中, 即: 默认的运算设备为 CPU。我们也可以将张量创建在 GPU 上, 能够利用对于矩阵计算的优势加快模型训练。将张量移动到 GPU 上有三种方法: 1. 使用 cuda 方法；；2. 直接在 GPU 上创建张量；3. 使用 to 方法指定设备

使用 cuda 方法

import torch

data = torch.tensor([10, 20 ,30])
print('存储设备:', data.device)

# 如果安装的不是 gpu 版本的 PyTorch，或电脑本身没有 NVIDIA 卡的计算环境，下面代码可能会报错
data = data.cuda()
print('存储设备:', data.device)

# 使用 cpu 函数将张量移动到 cpu 上
data = data.cpu()
print('存储设备:', data.device)

直接在 GPU 上创建张量

import torch

data = torch.tensor([10, 20, 30], device='cuda:0')
print('存储设备:', data.device)

使用 to 方法指定设备

import torch

data = torch.tensor([10, 20, 30])
print('存储设备:', data.device)

data = data.to('cuda:0')
print('存储设备:', data.device)

data = data.to('cpu')
print('存储设备:', data.device)

存储在不同设备的张量不能运算

import torch

data1 = torch.tensor([10, 20, 30], device='cuda:0')
data2 = torch.tensor([10, 20, 30])
print(data1.device, data2.device)

# RuntimeError: Expected all tensors to be on the same device,
# but found at least two devices, cuda:0 and cpu!
data = data1 + data2
print(data)

三、张量与numpy转换

1. 张量转换为 numpy 数组

使用 Tensor.numpy 函数可以将张量转换为 ndarray 数组，但是共享内存，可以使用 copy 函数避免共享。

import torch

data_tensor = torch.tensor([2, 3, 4])
# 使用张量对象中的 numpy 函数进行转换
data_numpy = data_tensor.numpy()
print(type(data_tensor))
print(type(data_numpy))

# 注意: data_tensor 和 data_numpy 共享内存，修改其中的一个，另外一个也会发生改变
# data_tensor[0] = 100
data_numpy[0] = 100
print(data_tensor)
print(data_numpy)

可以使用拷贝函数 copy() 产生新的数据，避免共享内存

data_numpy = data_tensor.numpy().copy()

2. numpy 数组转换为张量

# 默认进行的是浅拷贝
data_tensor = torch.from_numpy(data_numpy)

# 避免共享内存
data_tensor = torch.tensor(data_numpy)
data_tensor = torch.from_numpy(data_numpy.copy())

3. 提取单个元素的张量，使用 item

t1 = torch.tensor(30)
t2 = torch.tensor([30])
t3 = torch.tensor([[30]])

print(t1.shape)
print(t2.shape)
print(t3.shape)

print(t1.item())
print(t2.item())
print(t3.item())

# 注意: 张量中只有一个元素，如果有多个元素的话，使用 item 函数可能就会报错
# ValueError: only one element tensors can be converted to Python scalars
# t4 = torch.tensor([30, 40])
# print(t4.item())

四、张量拼接操作

## 1. cat 函数拼接，按照指定的维度进行拼接。

# 1. 按照0维度进行拼接
new_data = torch.cat([data1, data2], dim=0)
print(new_data.shape)

# 2. 按照1维度进行拼接
new_data = torch.cat([data1, data2], dim=1)
print(new_data.shape)

# 注意: dim 必须保证是有效的
# new_data = torch.cat([data1, data2], dim=3)
# print(new_data.shape)

2. Stack 函数的拼接，按照指定的维度进行叠加，会增加一个维度

data1 = torch.randint(0, 10, [2, 3])
data2 = torch.randint(0, 10, [2, 3])

# 将两个张量 stack 起来，像 cat 一样指定维度
# 1. 多出一个0维度，进行叠加
new_data = torch.stack([data1, data2], dim=0)
print(new_data.shape)

# 3. 多出一个2维度，进行叠加
new_data = torch.stack([data1, data2], dim=2)
print(new_data.shape)

五、张量的索引操作

1. 简单行列索引、列表索引

简单行列索引

# 1.1 获得指定的某行元素
print(data[2])

# 1.2 获得指定的某个列的元素
print(data[:,2])

# 冒号表示所有行或者所有列
print(data[:, :])

# 获得指定位置的某个元素：1行2列
print(data[1, 2], data[1][2])

# 表示先获得前三行，然后再获2列的数据
print(data[:3, 2])

# 表示获得前三行的前两列
print(data[:3, :2])

列表索引

# 如果索引的行列都是一个1维的列表，那么两个列表的长度必须相等
# 表示获得 (0, 0)、(2, 1)、(3, 2) 三个位置的元素
print(data[[0, 2, 3], [0, 1, 2]])

# 表示获得 0、2、3 行的 0、1、2 列
print(data[[[0], [2], [3]], [0, 1, 2]])

2. 布尔索引

# 希望能够获得该张量中所有大于3的元素
print(data[data > 3])

# 希望返回第2列元素大于6的行
print(data[data[:, 1] > 6])

# 希望返回第2行元素大于3的所有列
print(data[:, data[1] > 3])

3. 多维索引

# 按照第0个维度选择第0元素
print(data[0, :, :])

# 按照第1个维度选择第0元素
print(data[:, 0, :])

# 按照第2个维度选择第0元素
print(data[:, :, 0])

六、张量的形状操作

1. reshape 函数使用

data = torch.randint(0, 10, [4, 5])

# 查看张量的形状
print(data.shape, data.shape[0], data.shape[1])
print(data.size(), data.size(0), data.size(1))

# 修改张量的形状
new_data = data.reshape(2, 10)

# 注意: 转换之后的形状元素个数得等于原来张量的元素个数，不然会报错
# new_data = data.reshape(1, 10)
# print(new_data)

# 使用-1代替省略的形状，自动计算
new_data = data.reshape(5, -1)
new_data = data.reshape(-1, 2)

2. transpose 和 permute 函数的使用

transpose 函数可以一次性交换两个维度

# 将0列和2列交换
new_data = torch.transpose(data, 0, 2)

permute 函数可以一次性交换多个维度

# permute 函数可以一次性交换多个维度0,1,2变成1,2，0
new_data = torch.permute(data, [1, 2, 0])

3. view 函数也可以用于修改张量的形状，但是其用法比较局限，只能用于存储在`整块内存`中的张量。

当张量经过 transpose 或者 permute 函数之后，内存空间基本不连续,此时，必须先把空间连续，才能够使用 view 函数进行张量形状操作

data = torch.tensor([[10, 20, 30], [40, 50, 60]])
data = data.view(3, 2)

print('是否连续:', data.is_contiguous())
data = torch.transpose(data, 0, 1)
print('是否连续:', data.is_contiguous())

# 此时，在不连续内存的情况使用 view 会怎么样呢？需要先使用contiguous函数
data = data.contiguous().view(2, 3)
print(data)

4. squeeze 和 unsqueeze 函数使用

squeeze 函数用删除 shape 为 1 的维度

data = torch.randint(0, 10, [1, 3, 1, 5])
print(data.shape)

# 维度压缩, 默认去掉所有的1的维度
new_data = data.squeeze()
print(new_data.shape)

# 指定去掉某个1的维度
new_data = data.squeeze(2)
print(new_data.shape)

unsqueeze 可以添加一个为1的维度, 以增加数据的形状。

data = torch.randint(0, 10, [3, 5])
print(data.shape)

# 可以在指定位置增加维度
# -1 代表在最后一个维度上添加
new_data = data.unsqueeze(-1)
print(new_data.shape)

七、张量运算函数

import torch

torch.manual_seed(0)
# data = torch.randint(0, 10, [2, 3], dtype=torch.float64)
data = torch.randint(0, 10, [2, 3]).double()
print(data.dtype)

print(data)
# 默认对所有的数据计算均值
print(data.mean())
# 按指定的维度计算均值
print(data.mean(dim=0))
print(data.mean(dim=1))

# 2. 求和
print(data.sum())
print(data.sum(dim=0))
print(data.sum(dim=1))

# 3. 平方
print(data.pow(2))

# 4. 平方根
print(data.sqrt())

# 5. e多少次方
print(data.exp())

print(data.log())  # 以e为底
print(data.log2())  # 以2为底
print(data.log())  # 以10为底

八、模型保存加载

1. 直接序列化模型对象

直接使用 Python 的 pickle 模块将对象序列化到磁盘，该方法依赖于 PyTorch 的实现，可能受到 PyTorch 版本的影响。

import torch
import torch.nn as nn
import pickle

class Model(nn.Module):
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, input_size * 2)
        self.linear2 = nn.Linear(input_size * 2,  output_size)

    def forward(self, inputs):
        inputs = self.linear1(inputs)
        output = self.linear2(inputs)
        return output

# 保存模型
def test01():
    model = Model(128, 10)
    # 当我们的模型类继承了 nn.Module, 并且实现了 forward 函数
    # 此时，我们就可以把模型对象当做函数使用
    torch.save(model, 'model/test_model_save.pth', pickle_module=pickle, pickle_protocol=2)

# 加载模型
def test02():
    model = torch.load('model/test_model_save.pth',  pickle_module=pickle, map_location='cpu')
    print(model)

if __name__ == '__main__':
    test01()
    test02()

2. 存储模型参数

存储模型参数与 PyTorch 的实现关系较弱，建议使用第二种方法来存储模型。

import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, input_size * 2)
        self.linear2 = nn.Linear(input_size * 2,  output_size)

    def forward(self, inputs):
        inputs = self.linear1(inputs)
        output = self.linear2(inputs)
        return output

# 1. 实现模型参数存储
def test01():
    # 初始化模型参数
    model = Model(128, 10)
    # 初始化优化器
    optimizer = optim.Adam(model.parameters(), lr=1e-2)
    # 定义要存储的模型参数
    save_params = {
        'init_params': {'input_size': 128, 'output_size': 10},
        'acc_score': 0.98,
        'avg_loss': 0.86,
        'iter_num': 100,
        'optim_params': optimizer.state_dict(),
        'model_params': model.state_dict()
    }
    # 存储模型参数
    torch.save(save_params, 'model/model_params.pth')

# 2. 实现模型参数加载
def test02():
    # 从磁盘中将参数加载到内存中
    model_params = torch.load('model/model_params.pth')
    # 使用参数初始化模型
    model = Model(model_params['init_params']['input_size'], model_params['init_params']['output_size'])
    model.load_state_dict(model_params['model_params'])
    # 使用参数初始化优化器
    optimizer = optim.Adam(model.parameters())
    optimizer.load_state_dict(model_params['optim_params'])
    # 可以加载其他参数
    print('迭代次数:', model_params['iter_num'])
    print('准确率:', model_params['acc_score'])
    print('平均损失:', model_params['avg_loss'])

if __name__ == '__main__':
    test01()
    test02()