# 1.导入依赖包
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from torchsummary import summary
# 2.构建数据集
def create_dataset():
# 2.1 读取数据集
data = pd.read_csv('dataset/手机价格预测.csv')
# 2.2 获取特征值和目标值,类型转化 特征(Float) 标签(Long)
x, y = data.iloc[:, :-1], data.iloc[:, -1]
x, y = x.astype(np.float32), y.astype(np.int64)
# 2.3 数据集划分
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2,random_state=2)
# 2.4 数据转Tensor
train_dataset = TensorDataset(torch.from_numpy(x_train.values), torch.tensor(y_train.values))
test_dataset = TensorDataset(torch.from_numpy(x_test.values), torch.tensor(y_test.values))
return train_dataset, test_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__()
self.linear1 = nn.Linear(input_dim, 256)
self.linear2 = nn.Linear(256, 1024)
self.fc = nn.Linear(1024, output_dim)
def forward(self, x):
x = torch.relu(self.linear1(x))
x = torch.relu(self.linear2(x))
output = self.fc(x)
# output = torch.softmax(self.fc(x), dim=-1)
return output
# 4.模型训练(225)
def train(model, train_dataset, num_epochs, batch_size):
# 2 初始化参数 损失函数 优化器
loss1 = nn.CrossEntropyLoss()
# optimizer = optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
optimizer = optim.Adam(model.parameters(), lr=1e-4, betas=(0.99, 0.99))
start = time.time()
# 2 2个遍历 epoch dataloader
for epoch in range(num_epochs):
dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
total_num = 0
total_loss = 0.0
for x, y in dataloader:
# 5 前向传播 损失计算 梯度归零 反向传播 参数更新
output = model(x)
loss = loss1(output, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_num += 1 # 批次
total_loss += loss.item()
epoch += 1
print(f'epoch:{epoch + 1:4d},loss:{total_loss / (total_num * epoch):.4f}, time:{time.time() - start:.2f}s')
# 模型持久化
torch.save(model.state_dict(), 'model/phone2.pth')
# 5.模型预测评估
def test(model, test_dataset, input_dim, output_dim):
# 3.导入数据
dataloader = DataLoader(test_dataset, batch_size=8, shuffle=False)
correct = 0
# 4.遍历数据
for x, y in dataloader:
# 4.1 前向传播
output = model(x)
print(output)
# 4.2 获取输出结果(类别)
y_pred = torch.argmax(output, dim=1)
# print(y_pred) # 预测错误
# 4.3 计算准确率Acc
correct += (y_pred == y).sum()
print(correct.item())
Acc = correct.item() / len(test_dataset)
return Acc
if __name__ == '__main__':
train_dataset, test_dataset, feature_num, label_num = create_dataset()
# 1.实例化模型
model = PhonePriceModel(feature_num, label_num)
# 2.加载模型
model.load_state_dict(torch.load('model/phone2.pth'))
# 模型训练
# train(model, train_dataset, num_epochs=50, batch_size=8)
# 模型预测
Acc = test(model, test_dataset, feature_num, label_num)
print(f'Acc:{Acc:.5f}')