自定义数据集 使用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()}')
相关推荐
深蓝电商API几秒前
Firecrawl、ScrapeGraphAI:AI爬虫框架横评
人工智能·爬虫·firecrawl·scrapegraphai
Hpc_Fun3 分钟前
医学影像 AI 工作站怎么配?五类工位、3 秒首幅的时延预算与验收脚本
人工智能
薛定猫AI11 分钟前
【技术干货】多AI团队大模型协作编排与编程项目评测架构解析
python·机器学习·ai·大模型·技术实战
FII工业富联科技服务14 分钟前
如何提升工厂空间利用率与产线 UPH?从 Omniverse 布局优化到柔性自动化改造
人工智能·机器人·自动化·智能制造·omniverse·柔性生产·ai factory
深海鱼肝油ya19 分钟前
企业级RAG问答系统(一)数据处理模块
人工智能·大模型·pdf文件解析·rag问答系统·智能分块
本原财经22 分钟前
大厂争霸AI办公:个人Agent和AI硬件哪个才是答案?
大数据·人工智能
博谷26 分钟前
2026GEO概率化观测怎么做?BUGOOAI布谷GEO的采样与指标体系拆解
人工智能·ai
强哥聊算力31 分钟前
《算力江湖·国产GPU突围:摩尔线程故事》第四集:从图形到AI:摩尔线程为什么必须冲进大模型赛道?
运维·服务器·人工智能·ai·gpu算力
点PY41 分钟前
OmDet-Turbo: Real-Time Transformer-Based Open-Vocabulary Object Detection论文精读
深度学习·目标检测·transformer
空堂与归1 小时前
百度开源 LoopVL:1B 循环视觉模型省 260 倍数据
人工智能·百度·ai·开源