3. PyTorch基础回顾：张量操作、Dataset与DataLoader、模型构建三要素、训练循环模板

各位同学，欢迎来到第三章。在正式进入Transformer之前，咱们得先把PyTorch这把刀磨快。说实话，我见过太多同学模型结构背得滚瓜烂熟，结果一写训练循环就报错，一调张量维度就懵圈。这一章，咱们就把这些基本功砸实了。

3.1 张量操作：时间序列的基石

张量，说白了就是PyTorch里的数组。但时间序列预测里，张量的维度操作特别容易翻车。我个人习惯把张量想象成一个「多层抽屉的柜子」——每个维度就是一层抽屉的编号。

3.1.1 创建与基本属性

import torch

# 创建一个3x4的张量
x = torch.randn(3, 4)
print(x.shape)  # torch.Size([3, 4])
print(x.ndim)   # 2

# 时间序列常用：创建连续序列
seq = torch.arange(0, 100, step=1)  # 0到99
print(seq.shape)  # torch.Size([100])

嗯，这里要注意：时间序列数据通常是三维的——(batch_size, seq_len, input_dim)。我在做电力负荷预测时，就吃过这个亏。一开始把数据搞成了(seq_len, batch_size, input_dim)，结果模型训练了三天，精度一塌糊涂。后来才发现是维度顺序搞反了。

3.1.2 维度变换：最常用的三板斧

操作	函数	说明
重塑形状	view() / reshape()	不改变数据，只改变视图
增加维度	unsqueeze()	在指定位置插入一个维度
交换维度	transpose() / permute()	调换两个或多个维度的顺序

# 实战场景：把一维序列变成三维输入
seq = torch.randn(100)  # 原始序列
# 变成 (batch=1, seq_len=100, input_dim=1)
seq_3d = seq.unsqueeze(0).unsqueeze(-1)
print(seq_3d.shape)  # torch.Size([1, 100, 1])

# 交换维度：Transformer里经常用
x = torch.randn(2, 3, 4)  # (batch, seq, feature)
x_permuted = x.permute(0, 2, 1)  # (batch, feature, seq)
print(x_permuted.shape)  # torch.Size([2, 4, 3])

3.2 Dataset与DataLoader：数据管线的艺术

做时间序列预测，数据预处理占了60%的工作量。PyTorch的Dataset和DataLoader就是帮你把这60%的活干利索的。

3.2.1 自定义Dataset

我建议你养成一个习惯：所有时间序列数据都写成Dataset类。为什么？因为滑动窗口、归一化、缺失值处理，这些逻辑封装在一起，后期调试起来特别爽。

from torch.utils.data import Dataset, DataLoader

class TimeSeriesDataset(Dataset):
    def __init__(self, data, seq_len=24, pred_len=1):
        self.data = torch.FloatTensor(data)
        self.seq_len = seq_len
        self.pred_len = pred_len
    
    def __len__(self):
        # 总共有多少个样本
        return len(self.data) - self.seq_len - self.pred_len + 1
    
    def __getitem__(self, idx):
        # 返回 (输入序列, 目标序列)
        x = self.data[idx:idx + self.seq_len]
        y = self.data[idx + self.seq_len:idx + self.seq_len + self.pred_len]
        return x, y

# 使用示例
raw_data = [i for i in range(1000)]
dataset = TimeSeriesDataset(raw_data, seq_len=24, pred_len=1)
print(len(dataset))  # 976个样本

3.2.2 DataLoader的实用技巧

DataLoader有几个参数，我每次写代码都会仔细检查：

• batch_size：别贪大，显存不够就调小。我一般从32开始试。
• shuffle：训练集用True，验证集和测试集用False。时间序列里尤其要注意——打乱顺序会破坏时间依赖关系。
• num_workers：数据加载的并行进程数。我习惯设成4或8，但Windows上设成0更稳。

dataloader = DataLoader(
    dataset, 
    batch_size=32, 
    shuffle=True, 
    num_workers=4,
    drop_last=True  # 丢掉最后一个不完整的batch
)

for batch_x, batch_y in dataloader:
    print(batch_x.shape)  # torch.Size([32, 24, 1])
    print(batch_y.shape)  # torch.Size([32, 1, 1])
    break

3.3 模型构建三要素

构建一个PyTorch模型，说白了就三件事：继承nn.Module、定义层、实现forward。我见过有人把这三件事搞混，结果模型跑起来各种报错。

3.3.1 三要素详解

import torch.nn as nn

class SimpleTransformer(nn.Module):
    # 要素一：继承nn.Module
    def __init__(self, input_dim, d_model, nhead):
        # 要素二：定义层
        super().__init__()
        self.embedding = nn.Linear(input_dim, d_model)
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, 
            nhead=nhead,
            batch_first=True  # 这个参数很重要！
        )
        self.output = nn.Linear(d_model, 1)
    
    # 要素三：实现forward
    def forward(self, x):
        # x shape: (batch, seq_len, input_dim)
        x = self.embedding(x)
        x = self.encoder_layer(x)
        x = self.output(x[:, -1, :])  # 取最后一个时间步
        return x

3.3.2 参数初始化

模型初始化这事儿，很多人觉得无所谓。但我告诉你，初始化不好，模型可能根本不收敛。我一般会这样写：

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)

model = SimpleTransformer(input_dim=1, d_model=64, nhead=4)
model.apply(init_weights)

3.4 训练循环模板

训练循环，说白了就是「前向传播 → 算损失 → 反向传播 → 更新参数」这四个步骤的循环。我写过一个通用模板，用了好几年，改改就能用。

3.4.1 标准训练模板

def train_one_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    
    for batch_x, batch_y in dataloader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        
        # 1. 前向传播
        pred = model(batch_x)
        
        # 2. 算损失
        loss = criterion(pred, batch_y)
        
        # 3. 反向传播
        optimizer.zero_grad()
        loss.backward()
        
        # 4. 更新参数
        optimizer.step()
        
        total_loss += loss.item()
    
    return total_loss / len(dataloader)

def validate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0
    
    with torch.no_grad():
        for batch_x, batch_y in dataloader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            pred = model(batch_x)
            loss = criterion(pred, batch_y)
            total_loss += loss.item()
    
    return total_loss / len(dataloader)

3.4.2 完整训练流程

# 配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleTransformer(input_dim=1, d_model=64, nhead=4).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

# 训练
num_epochs = 100
best_loss = float('inf')

for epoch in range(num_epochs):
    train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
    val_loss = validate(model, val_loader, criterion, device)
    
    # 保存最佳模型
    if val_loss < best_loss:
        best_loss = val_loss
        torch.save(model.state_dict(), 'best_model.pth')
    
    if (epoch + 1) % 10 == 0:
        print(f'Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}')

3.5 本章知识体系

下面这张图，是我自己总结的PyTorch时间序列预测的核心流程。你把它存下来，写代码时对照着看，能少走很多弯路。

嗯，这一章的内容就到这里。这些基础操作，你写代码时肯定会反复用到。别嫌烦，多练几遍，形成肌肉记忆就好了。

---

公众号:蓝海资料掘金营，微信deep3321