自定义数据集 使用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()}')
相关推荐
jason_yang5 小时前
写给女儿的英语App
前端·人工智能
2503_931712485 小时前
2026 企业级 AI 智能体平台选型指南:市场格局、品类解析与落地建议
大数据·人工智能
创世宇图5 小时前
从 8% 到 61%,8 倍提升!Meta 开源 Brain2Qwerty v2:非侵入式脑机接口首次逼近侵入式水平
人工智能
LadenKiller5 小时前
近期AI量化辅助,工具重点要跟学习阶段变化
人工智能·python
SimpleLearingAI6 小时前
DiT:Diffusion Transformer原理简介
人工智能·深度学习·transformer
满怀冰雪6 小时前
08-Paddle 神经网络层入门:Linear、激活函数与 Sequential
神经网络·机器学习·paddle
幸福在路上wellbeing6 小时前
AI 智能体开发第一月 · Week 1 逐小时执行手册
人工智能
tedcloud1236 小时前
Orca 部署指南:开源 AI 推理服务的 Linux 部署实践
linux·运维·服务器·人工智能·开源·自动化·excel
AcaDesign6 小时前
上海市东方英才计划项目答辩PPT设计案例模板 | WordinPPT
大数据·人工智能
搭贝6 小时前
企业智能知识库搭建实战:搭贝低代码平台实现AI知识管理系统
人工智能·低代码