# 1.导入相关模块
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import timefrom torchsummary import summary# 2.构建数据集
def create_dataset():# 使用pandas读取数据data = pd.read_csv('dataset/手机价格预测.csv')# 特征值和目标值x, y = data.iloc[:, :-1], data.iloc[:, -1]# 类型转换:特征值,目标值x = x.astype(np.float32)y = y.astype(np.int64)# 数据集划分x_train, x_valid, y_train, y_valid = train_test_split(x, y, train_size=0.8, random_state=88)# 构建数据集,转换为pytorch的形式train_dataset = TensorDataset(torch.from_numpy(x_train.values), torch.tensor(y_train.values))valid_dataset = TensorDataset(torch.from_numpy(x_valid.values), torch.tensor(y_valid.values))# 返回结果return train_dataset, valid_dataset, x_train.shape[1], len(np.unique(y))# 3.构建网络模型
class PhonePriceModel(nn.Module):def __init__(self, input_dim, output_dim):super(PhonePriceModel, self).__init__()# 1. 第一层:输入维度:20,输出维度:128self.linear1 = nn.Linear(input_dim, 128)# 2. 第二层:输入维度:128,输出维度:256self.linear2 = nn.Linear(128, 256)# 3. 第三层:输入维度:256,输出维度:4self.linear3 = nn.Linear(256, output_dim)def forward(self, x):# 前向传播过程x = torch.relu(self.linear1(x))x = torch.relu(self.linear2(x))output = self.linear3(x)# 获取数据结果return output# 4.模型训练
def train(train_dataset, input_dim, class_num, ):# 固定随机数种子torch.manual_seed(0)# 初始化模型model = PhonePriceModel(input_dim, class_num)# 损失函数criterion = nn.CrossEntropyLoss()# 优化方法optimizer = optim.SGD(model.parameters(), lr=1e-3)# 训练轮数num_epoch = 50# 遍历每个轮次的数据for epoch_idx in range(num_epoch):# 初始化数据加载器dataloader = DataLoader(train_dataset, shuffle=True, batch_size=8)# 训练时间start = time.time()# 计算损失total_loss = 0.0total_num = 1# 遍历每个batch数据进行处理for x, y in dataloader:# 将数据送入网络中进行预测output = model(x)# 计算损失loss = criterion(output, y)# 梯度归零optimizer.zero_grad()# 反向传播loss.backward()# 参数更新optimizer.step()# 损失计算total_num += 1total_loss += loss.item()# 打印损失变换结果print('epoch: %4s loss: %.2f, time: %.2fs' % (epoch_idx + 1, total_loss / total_num, time.time() - start))# 模型保存torch.save(model.state_dict(), 'model/phone.pth')def test(valid_dataset, input_dim, class_num):# 加载模型和训练好的网络参数model = PhonePriceModel(input_dim, class_num)model.load_state_dict(torch.load('model/phone.pth'))# 构建加载器dataloader = DataLoader(valid_dataset, batch_size=8, shuffle=False)# 评估测试集correct = 0# 遍历测试集中的数据for x, y in dataloader:# 将其送入网络中output = model(x)# 获取类别结果y_pred = torch.argmax(output, dim=1)# 获取预测正确的个数correct += (y_pred == y).sum()# 求预测精度print('Acc: %.5f' % (correct.item() / len(valid_dataset)))if __name__ == '__main__':# 1.获取数据train_dataset, valid_dataset, input_dim, class_num = create_dataset()print("输入特征数:", input_dim)print("分类个数:", class_num)# 2.模型实例化model = PhonePriceModel(input_dim, class_num)summary(model, input_size=(input_dim,), batch_size=16)# 3.模型训练# train(train_dataset, input_dim, class_num)# 4.模型预测test(valid_dataset, input_dim, class_num)
优化点: