import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

一、输入数据

# 散点输入数据  
data = np.array([  
    [-0.5, 7.7], [1.8, 98.5], [0.9, 57.8], [0.4, 39.2],  
    [-1.4, -15.7], [-1.4, -37.3], [-1.8, -49.1], [1.5, 75.6],  
    [0.4, 34.0], [0.8, 62.3]  
])  

# 提取x和y并转换为tensor  
x_train = torch.tensor(data[:, 0], dtype=torch.float32)  
y_train = torch.tensor(data[:, 1], dtype=torch.float32)  

# 转换为 TensorDataset  
dataset = TensorDataset(x_train, y_train)  

# 如需创建 DataLoader，可调整批量大小  
batch_size = 2  # 可根据需求进行调整  
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)  


# 设置随机数种子
seed = 42
torch.manual_seed(seed)

二、定义前向模型

class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        x = self.linear(x)
        return x


model = LinearModel()
#定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

三、训练模型

# 创建 DataLoader
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

for epoch in range(1, 501):
    total_loss = 0
    for batch_x, batch_y in dataloader:
        # 前向传播
        y_hat = model(batch_x.unsqueeze(1))
        # 计算损失
        loss = criterion(y_hat.squeeze(1), batch_y)
        total_loss += loss.item()

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 计算平均损失
    avg_loss = total_loss / len(dataloader)

    # 打印损失
    if epoch % 10 == 0 or epoch == 1:
        print(f"Epoch: {epoch}, Loss: {avg_loss}")

四、模型的保存与加载+预测

4.1、整个模型保存

        将整个模型保存到名为 'entire_model.pth' 的文件中。包括模型的结构和参数。文件中将包含模型的所有信息，包括网络结构、权重等。 但是要注意，保存整个模型可能会占用更多 的磁盘空间，并且不如保存状态字典灵活， 因为状态字典可以与不同的模型结构兼容。

torch.save(model, 'entire_model.pth')

 模型加载

model = torch.load('./entire_model.pth')
# 评估模型
model.eval()

注意：

        要将模型设置为评估模式，在评估模式下，不需要计算梯度

模型预测 

x_test = torch.tensor([[1.2], [3.4], [5.6]], dtype=torch.float32)
with torch.no_grad():
    y_pred = model(x_test)
print(y_pred)
'''
tensor([[ 70.6924],
        [155.1746],
        [239.6569]])
'''

4.2、只保存模型的参数

        model.state_dict() 返回模型的状态字 典，其中包含模型的所有参数。 然后，torch.save() 函数将这个状态字 典保存到名为 'model.pth' 的文件中。

        只保存模型的参数（即模型的状态字典），不保存模型的结构。文件中将只包含模型的参数，没有模型的结构信息。

torch.save(model.state_dict(), 'model.pth')

 模型加载 

model.load_state_dict(torch.load('model.pth'))
# 评估模型
model.eval()

注意：

        要将模型设置为评估模式，在评估模式下，不需要计算梯度

 模型预测 

x_test = torch.tensor([[1.2], [3.4], [5.6]], dtype=torch.float32)
with torch.no_grad():  
    y_pred = model(x_test)
print(y_pred)
'''
tensor([[ 70.6924],
        [155.1746],
        [239.6569]])
'''

五、完整代码

5.1、保存模型

import torch.nn as nn  
import numpy as np  
import torch  
from torch.utils.data import DataLoader, TensorDataset  

# 1. 散点输入数据  
data = np.array([  
    [-0.5, 7.7], [1.8, 98.5], [0.9, 57.8], [0.4, 39.2],  
    [-1.4, -15.7], [-1.4, -37.3], [-1.8, -49.1], [1.5, 75.6],  
    [0.4, 34.0], [0.8, 62.3]  
])  

# 提取x和y并转换为tensor  
x_train = torch.tensor(data[:, 0], dtype=torch.float32)  # 提取特征（X）  
y_train = torch.tensor(data[:, 1], dtype=torch.float32)  # 提取标签（Y）  

# 转换为 TensorDataset  
dataset = TensorDataset(x_train, y_train)  # 将特征和标签组合成数据集  

# 如需创建 DataLoader，可调整批量大小  
batch_size = 2  # 可根据需求进行调整  
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)  # 创建数据加载器  

# 设置随机数种子  
seed = 42  
torch.manual_seed(seed)  # 为可重复性设置随机种子  

# 2、定义前向模型  
class LinearModel(nn.Module):  
    def __init__(self):  
        super(LinearModel, self).__init__()  # 调用父类构造函数  
        self.linear = nn.Linear(1, 1)  # 定义线性层，输入和输出都是1维  

    def forward(self, x):  
        x = self.linear(x)  # 前向传播  
        return x  

model = LinearModel()  # 实例化线性模型  

# 3、定义损失函数和优化器  
criterion = nn.MSELoss()  # 使用均方误差作为损失函数  
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 使用随机梯度下降作为优化器  

# 创建 DataLoader  
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)  # 重新定义数据加载器 | 可选：使用不同的batch size  

# 进行模型训练  
for epoch in range(1, 501):  # 训练500个epoch  
    total_loss = 0  # 初始化总损失  
    for batch_x, batch_y in dataloader:  # 遍历数据加载器中的每个批次  
        # 前向传播  
        y_hat = model(batch_x.unsqueeze(1))  # 将特征添加一个维度以匹配模型输入  
        # 计算损失  
        loss = criterion(y_hat.squeeze(1), batch_y)  # 计算当前批次的损失  
        total_loss += loss.item()  # 记录损失  

        # 反向传播和优化  
        optimizer.zero_grad()  # 清空之前的梯度  
        loss.backward()  # 反向传播，计算梯度  
        optimizer.step()  # 更新模型参数  

    # 计算平均损失  
    avg_loss = total_loss / len(dataloader)  # 计算当前epoch的平均损失  

# 保存模型  
# 1)这个方法将整个模型保存到文件中，包括模型的结构和参数。  
# 文件中将包含模型的所有信息，包括网络结构、权重等。  
torch.save(model, 'entire_model.pth')  

# 2)这个方法只保存模型的参数（即模型的状态字典），不保存模型的结构。  
# 文件中将只包含模型的参数，没有模型的结构信息。  
torch.save(model.state_dict(), 'model.pth')  

5.2、加载模型预测

import torch
from torch import nn

# 定义线性模型类
class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()  # 调用父类构造函数
        self.linear = nn.Linear(1, 1)  # 定义一个线性层，输入和输出都是1维

    def forward(self, x):
        x = self.linear(x)  # 前向传播
        return x

# 加载模型
# 如果保存的是整个模型，使用下面这行代码
model = torch.load('entire_model.pth',weights_only=False)    # 加载完整模型（包括结构和参数）

# 如果保存的是模型的状态字典，使用下面这行代码
# model.load_state_dict(torch.load('model.pth'))  # 仅加载模型参数

# 评估模型
model.eval()  # 设置模型为评估模式，关闭dropout和batch normalization等

# 使用模型进行预测
x_test = torch.tensor([[1.2], [3.4], [5.6]], dtype=torch.float32)  # 创建测试数据
with torch.no_grad():  # 在评估时不计算梯度以减少内存使用和提高速度
    y_pred = model(x_test)  # 模型预测

print(y_pred)  # 输出预测值
'''
tensor([[ 70.6924],
        [155.1746],
        [239.6569]])
'''

六、库函数

6.1、TensorDataset()

        是 PyTorch 中的一个类，用于将多个张量组合成一个数据集。它方便地将输入数据和目标数据成对存储，可以通过索引访问。

6.2、DataLoader()

        是一个 PyTorch 提供的工具，用于在训练过程中批量加载数据。它支持打乱数据、并行加载、以及自动生成批次。

DataLoader(dataset, batch_size=2, shuffle=True)

方法	描述
dataset	数据集
batch_size	每个批次的样本数量。
shuffle	是否在每个epoch前打乱数据。

6.3、MSELoss()

        是 PyTorch 中用于计算均方误差损失（Mean Squared Error Loss）的类。通常用于回归任务，衡量预测值和真实值之间的差异。

criterion(predictions, targets) 

方法	描述
predictions	是模型的预测值
targets	是真实值

6.4、SGD()

        代表随机梯度下降（Stochastic Gradient Descent），是用于优化模型的算法。它根据当前梯度来更新模型参数。

torch.optim.SGD(model.parameters(), lr=0.01)

方法	描述
params	要优化的参数（通常是模型的参数）。
lr	学习率（learning rate），决定每次参数更新的步长。

6.5、zero_grad()

        用来清空优化器中累积的梯度，以避免梯度在每次迭代中累加。

6.6、backward()

        用来计算当前损失相对于模型参数的梯度。

6.7、step()

        方法用于根据当前累积的梯度更新模型参数。