论文辅助笔记：TEMPO 之 dataset.py

0 导入库

import os
import pandas as pd
import torch
from torch.utils.data import Dataset
from .utils import StandardScaler, decompose
from .features import time_features

1 Dataset_ETT_hour

1.1 构造函数

class Dataset_ETT_hour(Dataset):def __init__(self,root_path,flag="train",size=None,features="S",data_path="ETTh1.csv",target="OT",scale=True,inverse=False,timeenc=0,freq="h",cols=None,period=24,):if size == None:self.seq_len = 24 * 4 * 4self.pred_len = 24 * 4else:self.seq_len = size[0]self.pred_len = size[1]#输入sequence和输出sequence的长度assert flag in ["train", "test", "val"]type_map = {"train": 0, "val": 1, "test": 2}self.set_type = type_map[flag]'''指定数据集的用途，可以是 "train"、"test" 或 "val"，分别对应训练集、测试集和验证集'''self.features = features#指定数据集包含的特征类型，默认为 "S"，表示单一特征self.target = target#指定预测的目标特征self.scale = scale#一个布尔值，用于确定数据是否需要归一化处理self.inverse = inverse#一个布尔值，用于决定是否进行逆变换self.timeenc = timeenc#用于确定是否对时间进行编码【原始模样 or -0.5~0.5区间】self.freq = freq#定义时间序列的频率，如 "h" 表示小时级别的频率self.period = period#定义时间序列的周期，默认为 24self.root_path = root_pathself.data_path = data_pathself.__read_data__()#用于读取并初始化数据集

1.2 __read_data__

def __read_data__(self):self.scaler = StandardScaler()#初始化一个 StandardScaler 对象，用于数据的标准化处理df_raw = pd.read_csv(os.path.join(self.root_path, self.data_path))#读取数据集文件，将其存储为 DataFrame 对象 df_rawborder1s = [0,12 * 30 * 24 - self.seq_len,12 * 30 * 24 + 4 * 30 * 24 - self.seq_len,]#定义了三个区间的起始位置，分别对应训练集、验证集和测试集border2s = [12 * 30 * 24,12 * 30 * 24 + 4 * 30 * 24,12 * 30 * 24 + 8 * 30 * 24,]#定义了每个区间的结束位置border1 = border1s[self.set_type]border2 = border2s[self.set_type]'''通过 self.set_type 确定当前数据集类型并从 border1s 和 border2s 中获取对应的起始和结束位置 border1 和 border2'''if self.features == "M" or self.features == "MS":cols_data = df_raw.columns[1:]df_data = df_raw[cols_data]elif self.features == "S":df_data = df_raw[[self.target]]'''选择特征数据：多特征 "M" 或 "MS"：选择所有数据列，除去日期列。单一特征 "S"：只选择目标特征列（由 self.target 指定）。'''if self.scale:train_data = df_data[border1s[0] : border2s[0]]self.scaler.fit(train_data.values)data = self.scaler.transform(df_data.values)else:data = df_data.values'''如果 self.scale 为 True，则执行数据归一化：train_data：选择训练集的数据，用于拟合 self.scaler。data：对整个 df_data 进行转换。'''df_stamp = df_raw[["date"]][border1:border2]df_stamp["date"] = pd.to_datetime(df_stamp.date)data_stamp = time_features(df_stamp, timeenc=self.timeenc, freq=self.freq)'''时间特征处理：提取日期列 df_stamp，并将其转换为时间特征：pd.to_datetime：将日期转换为 datetime 对象。time_features：用于生成时间特征。'''self.data_x = data[border1:border2]if self.inverse:self.data_y = df_data.values[border1:border2]else:self.data_y = data[border1:border2]self.data_stamp = data_stamp'''将转换后的数据和时间特征赋值给 self.data_x、self.data_y 和 self.data_stamp：self.data_x 取 data 中的对应区间数据。self.data_y 根据 self.inverse 决定是从 data 还是 df_data 中获取。self.data_stamp 取生成的时间特征。'''

1.3 getitem

def __getitem__(self, index):s_begin = index#设置序列的起始点s_end = s_begin + self.seq_len#计算序列的结束点r_begin = s_end#设置预测序列的起始点r_end = r_begin + self.pred_len#计算预测序列的结束点seq_x = self.data_x[s_begin:s_end]#从 data_x 中提取序列部分seq_y = self.data_y[r_begin:r_end]# 从 data_y 中提取预测部分[ground-truth]x = torch.tensor(seq_x, dtype=torch.float).transpose(1, 0)  # [1, seq_len]y = torch.tensor(seq_y, dtype=torch.float).transpose(1, 0)  # [1, pred_len](trend, seasonal, residual) = decompose(x, period=self.period)#对序列 x 进行时间序列分解，返回趋势、季节性和残差三部分components = torch.cat((trend, seasonal, residual), dim=0)  # [3, seq_len]#将分解后的三部分按 0 维（纵向）拼接，形成一个包含三种特征的张量return components, y

1.3len

    def __len__(self):return len(self.data_x) - self.seq_len - self.pred_len + 1

1.4 inverse_transform

将数据进行逆转换，还原到原始尺度

    def inverse_transform(self, data):return self.scaler.inverse_transform(data)

2 Dataset_ETT_minute

基本上和hour 的一样，几个地方不一样：

__init__
- data_path="ETTm1.csv",
- freq="t",
- period: int = 60,

__read_data__

border1s = [0,12 * 30 * 24 * 4 - self.seq_len,12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len,]
border2s = [12 * 30 * 24 * 4,12 * 30 * 24 * 4 + 4 * 30 * 24 * 4,12 * 30 * 24 * 4 + 8 * 30 * 24 * 4,]