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

发布于:2025-02-10 ⋅ 阅读:(99) ⋅ 点赞:(0)

import torch
import torch.nn as nn
from torch.utils.data import Dataset
import numpy as np

# 自定义数据集类
class CustomDataset(Dataset):
    def __init__(self, x_data, y_data):
        self.x_data = torch.from_numpy(x_data).float()
        self.y_data = torch.from_numpy(y_data).float()

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

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

# 逻辑回归模型
class LogisticRegressionModel(nn.Module):
    def __init__(self, input_dim):
        super(LogisticRegressionModel, self).__init__()
        self.linear = nn.Linear(input_dim, 1)  # 输出维度为1,因为这是二分类问题

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

# 创建数据集
x_data = np.array([[1], [2], [3], [4], [5]], dtype=np.float32)
y_data = np.array([[0], [0], [1], [1], [1]], dtype=np.float32)
dataset = CustomDataset(x_data, y_data)

# 创建数据加载器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)

# 创建模型、损失函数和优化器
model = LogisticRegressionModel(input_dim=1)
criterion = nn.BCELoss()  # 二分类交叉熵损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练模型
epochs = 100
for epoch in range(epochs):
    for x_batch, y_batch in dataloader:
        # 前向传播
        y_pred = model(x_batch)
        loss = criterion(y_pred, y_batch.view(-1, 1))
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

# 保存模型
torch.save(model.state_dict(), 'logistic_regression_model.pth')

# 加载模型
model = LogisticRegressionModel(input_dim=1)
model.load_state_dict(torch.load('logistic_regression_model.pth'))

# 进行预测
with torch.no_grad():
    x_test = torch.tensor([[6], [7], [8]], dtype=torch.float32)
    y_pred = model(x_test)
    print('预测值:', y_pred)
 


网站公告

今日签到

点亮在社区的每一天
去签到