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

python 复制代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import pandas as pd
import numpy as np
class CustomDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]
        if self.transform:
            sample = self.transform(sample)
        return sample, label
# 示例数据
data = np.random.rand(100, 10)  # 100个样本,每个样本10个特征
labels = np.random.randint(0, 2, 100)  # 二分类标签

# 转换为torch的Tensor
data = torch.tensor(data, dtype=torch.float32)
labels = torch.tensor(labels, dtype=torch.long)

# 创建数据集和数据加载器
dataset = CustomDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)


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

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


# 参数设置
input_dim = data.shape[1]
num_epochs = 20
learning_rate = 0.01

# 初始化模型、损失函数和优化器
model = LogisticRegression(input_dim)
criterion = nn.BCELoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        # 将标签转换为浮点型
        labels = labels.float().unsqueeze(1)

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

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

    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')
# 保存模型权重和架构
torch.save(model.state_dict(), 'logistic_regression_model.pth')

# 加载模型
loaded_model = LogisticRegression(input_dim)
loaded_model.load_state_dict(torch.load('logistic_regression_model.pth'))
loaded_model.eval()  # 切换为评估模式

# 进行预测
with torch.no_grad():
    sample = torch.tensor(np.random.rand(1, 10), dtype=torch.float32)  # 一个新的样本
    prediction = loaded_model(sample)
    print(f'Prediction: {prediction.item()}')
相关推荐
江畔柳前堤1 分钟前
github实战指南01-账号配置与 SSH 密钥
运维·人工智能·深度学习·ssh·github·pyqt·信号处理
workflower1 小时前
使用大语言模型处理用户需求
大数据·人工智能·设计模式·重构·动态规划
CodePlayer竟然被占用了2 小时前
没有生态的大模型不算前沿
人工智能
米小虾2 小时前
AI Agent 开发实战:2026年主流框架与MCP协议深度解析
人工智能·agent
米小虾2 小时前
2026年AI大模型半年报:从"参数军备"到"生态为王",谁在领跑下半场?
人工智能
m0_571186602 小时前
第五十周周报
人工智能
寰宇视讯2 小时前
解码AI未来 2026世界制造业大会人工智能与机器人展9月启幕
人工智能·机器人
冬奇Lab2 小时前
每日一个开源项目(第132篇):SkillSpector - 安装 AI Agent Skill 之前先扫一遍
人工智能·开源·agent
冬奇Lab2 小时前
如何让 AI Skill 质量有据可查?Benchmark 驱动的评测体系设计
人工智能·agent