ShopTRAINING/models/data_utils.py

50 lines
1.5 KiB
Python
Raw Normal View History

import torch
import numpy as np
from torch.utils.data import Dataset
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
# 定义数据集类
class PharmacyDataset(Dataset):
def __init__(self, data_X, data_Y):
self.data_X = data_X
self.data_Y = data_Y
def __getitem__(self, index):
return self.data_X[index], self.data_Y[index]
def __len__(self):
return len(self.data_X)
# 定义用于时间序列预测的数据处理函数
def create_dataset(datasetX, datasetY, look_back=1, predict_steps=1):
dataX, dataY = [], []
for i in range(0, len(datasetX) - look_back - predict_steps + 1):
x = datasetX[i:(i + look_back), :]
if predict_steps == 1:
y = datasetY[i + look_back]
else:
y = datasetY[i + look_back:i + look_back + predict_steps, 0] # 仅取销量列
dataX.append(x)
dataY.append(y)
return np.array(dataX), np.array(dataY)
# 评估函数
def evaluate_model(y_true, y_pred):
mse = mean_squared_error(y_true, y_pred)
rmse = np.sqrt(mse)
mae = mean_absolute_error(y_true, y_pred)
r2 = r2_score(y_true, y_pred)
# MAPE计算时避免除以0
mask = y_true != 0
y_true_masked = y_true[mask]
y_pred_masked = y_pred[mask]
mape = np.mean(np.abs((y_true_masked - y_pred_masked) / y_true_masked)) * 100
return {
'MSE': mse,
'RMSE': rmse,
'MAE': mae,
'': r2,
'MAPE(%)': mape
}