自定义数据集,使用 PyTorch 框架实现逻辑回归并保存模型,然后保存模型后再加载模型进行预测

在本文中,我们将展示如何使用 NumPy 创建自定义数据集,利用 PyTorch 实现一个简单的逻辑回归模型,并在训练完成后保存该模型,最后加载模型并用它进行预测。

1. 创建自定义数据集

首先,我们使用 NumPy 创建一个简单的二分类数据集。假设我们的数据集包含两个特征。

复制代码
import numpy as np

# 生成随机数据
np.random.seed(42)
X = np.random.randn(100, 2)  # 100个样本,2个特征
y = (X[:, 0] + X[:, 1] > 0).astype(int)  # 标签为1或0

# 打印数据集的前5个样本
print(X[:5], y[:5])

2. 构建逻辑回归模型

接下来,我们使用 PyTorch 来构建一个简单的逻辑回归模型。PyTorch 提供了 torch.nn.Module 类,能够轻松实现神经网络模型。

复制代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 转换 NumPy 数据为 PyTorch 张量
X_tensor = torch.tensor(X, dtype=torch.float32)
y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1)

# 创建数据集并加载
dataset = TensorDataset(X_tensor, y_tensor)
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

# 定义逻辑回归模型
class LogisticRegressionModel(nn.Module):
    def __init__(self):
        super(LogisticRegressionModel, self).__init__()
        self.linear = nn.Linear(2, 1)  # 2个输入特征,1个输出

    def forward(self, x):
        return torch.sigmoid(self.linear(x))

# 初始化模型
model = LogisticRegressionModel()

3. 训练模型

接下来,我们定义损失函数和优化器,并训练模型。

复制代码
# 定义损失函数和优化器
criterion = nn.BCELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
epochs = 1000
for epoch in range(epochs):
    for inputs, labels in dataloader:
        optimizer.zero_grad()  # 清空梯度
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新权重

    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

4. 保存模型

模型训练完成后,我们可以将模型保存到文件中。PyTorch 提供了 torch.save() 方法来保存模型的状态字典。

复制代码
# 保存模型
torch.save(model.state_dict(), 'logistic_regression_model.pth')
print("模型已保存!")

5. 加载模型并进行预测

我们可以加载保存的模型,并对新数据进行预测。

复制代码
# 加载模型
loaded_model = LogisticRegressionModel()
loaded_model.load_state_dict(torch.load('logistic_regression_model.pth'))
loaded_model.eval()  # 切换到评估模式

# 进行预测
with torch.no_grad():
    test_data = torch.tensor([[1.5, -0.5]], dtype=torch.float32)
    prediction = loaded_model(test_data)
    print(f'预测值: {prediction.item():.4f}')

6. 总结

在这篇博客中,我们展示了如何使用 NumPy 创建一个简单的自定义数据集,并使用 PyTorch 实现一个逻辑回归模型。我们还展示了如何保存训练好的模型,并加载模型进行预测。通过保存和加载模型,我们可以在不同的时间或环境中重复使用已经训练好的模型,而不需要重新训练它。

相关推荐
Ronin-Lotus3 小时前
深度学习篇---剪裁&缩放
图像处理·人工智能·缩放·剪裁
cpsvps4 小时前
3D芯片香港集成:技术突破与产业机遇全景分析
人工智能·3d
国科安芯4 小时前
抗辐照芯片在低轨卫星星座CAN总线通讯及供电系统的应用探讨
运维·网络·人工智能·单片机·自动化
AKAMAI4 小时前
利用DataStream和TrafficPeak实现大数据可观察性
人工智能·云原生·云计算
Ai墨芯1115 小时前
深度学习水论文:特征提取
人工智能·深度学习
无名工程师5 小时前
神经网络知识讨论
人工智能·神经网络
nbsaas-boot5 小时前
AI时代,我们更需要自己的开发方式与平台
人工智能
SHIPKING3935 小时前
【机器学习&深度学习】LLamaFactory微调效果与vllm部署效果不一致如何解决
人工智能·深度学习·机器学习
jonyleek6 小时前
如何搭建一套安全的,企业级本地AI专属知识库系统?从安装系统到构建知识体系,全流程!
人工智能·安全
MQ_SOFTWARE7 小时前
AI驱动的金融推理:Fin-R1模型如何重塑行业决策逻辑
人工智能·金融