论文辅助笔记：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 * 4
            self.pred_len = 24 * 4
        else:
            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
        #定义时间序列的周期，默认为 24

        self.root_path = root_path
        self.data_path = data_path

        self.__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_raw

        border1s = [
            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.3__len__

    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,
            ]