用 Python 轻松实现时间序列预测：Darts 使用 GPU 和 TPU 的 Torch 模型 Using Torch Models with GPUs and TPUs-CSDN博客

本文链接：https://blog.csdn.net/weixin_47339916/article/details/149747904

文中内容仅限技术学习与代码实践参考，市场存在不确定性，技术分析需谨慎验证，不构成任何投资建议。

Darts

Darts 是一个 Python 库，用于对时间序列进行用户友好型预测和异常检测。它包含多种模型，从 ARIMA 等经典模型到深度神经网络。所有预测模型都能以类似 scikit-learn 的方式使用 fit() 和 predict() 函数。该库还可以轻松地对模型进行回溯测试，将多个模型的预测结果结合起来，并将外部数据考虑在内。Darts 支持单变量和多变量时间序列和模型。基于 ML 的模型可以在包含多个时间序列的潜在大型数据集上进行训练，其中一些模型还为概率预测提供了丰富的支持。

使用 GPU 和 TPU 的 Torch 模型

Using Torch Models with GPUs and TPUs

本章节适用于 Darts 0.17.0 及更高版本。

我们假设您已经了解 Darts 中的 Torch 预测模型。如果您是新手，建议先阅读 Torch 预测模型指南。该指南还包含一个性能建议章节，建议优先阅读。此外，这里还有一个循环神经网络（RNN）模型示例，本章节将以此为基础展开。

使用 CPU

默认情况下，所有模型均在 CPU 上运行。如上方 RNN 示例所示，我们将导入 Air Passenger 数据集以及其他必要模块。

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

from darts.dataprocessing.transformers import Scaler
from darts.models import RNNModel
from darts.metrics import mape
from darts.datasets import AirPassengersDataset

接下来按如下方式读取并缩放数据：

# 读取数据：
series = AirPassengersDataset().load()
series = series.astype(np.float32)

# 创建训练集和验证集：
train, val = series.split_after(pd.Timestamp("19590101"))

# 归一化时间序列（注意：我们避免在验证集上拟合变换器）
transformer = Scaler()
train_transformed = transformer.fit_transform(train)
val_transformed = transformer.transform(val)
series_transformed = transformer.transform(series)

然后按如下方式创建 RNN：

my_model = RNNModel(
    model="RNN",
    hidden_dim=20,
    dropout=0,
    batch_size=16,
    n_epochs=300,
    optimizer_kwargs={"lr": 1e-3},
    model_name="Air_RNN",
    log_tensorboard=True,
    random_state=42,
    training_length=20,
    input_chunk_length=14,
    force_reset=True,
)

并拟合数据：

my_model.fit(train_transformed, val_series=val_transformed)

从输出可以看到，训练模型时未使用其他处理单元：

GPU available: False, used: False
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs

  | Name      | Type    | Params
--------------------------------------
0 | criterion | MSELoss | 0
1 | rnn       | RNN     | 460
2 | V         | Linear  | 21
--------------------------------------
481       Trainable params
0         Non-trainable params
481       Total params
0.004     Total estimated model params size (MB)

Epoch 299: 100% 8/8 [00:00<00:00, 42.49it/s, loss=0.00285, v_num=logs]
<darts.models.forecasting.rnn_model.RNNModel at 0x7fc75901cb10>

模型现已就绪，可以开始预测；因已在指南开头链接的示例中展示，此处不再重复。

使用 GPU

GPU 可以在处理时间方面显著提升模型性能。通过在 Pytorch Lightning Trainer 中使用 Accelerator，我们可以享受 GPU 带来的优势。只需通过 PyTorch Lightning Trainer 参数（以 pl_trainer_kwargs 字典形式）指示模型使用机器的 GPU，示例如下：

my_model = RNNModel(
    model="RNN",
    ...
    force_reset=True,
    pl_trainer_kwargs={
      "accelerator": "gpu",
      "devices": [0]
    },
)

此时输出如下：

GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name      | Type    | Params
--------------------------------------
0 | criterion | MSELoss | 0
1 | rnn       | RNN     | 460
2 | V         | Linear  | 21
--------------------------------------
481       Trainable params
0         Non-trainable params
481       Total params
0.004     Total estimated model params size (MB)

Epoch 299: 100% 8/8 [00:00<00:00, 39.81it/s, loss=0.00285, v_num=logs]
<darts.models.forecasting.rnn_model.RNNModel at 0x7ff1b5e4d4d0>

由输出可见，GPU 已可用并被使用。其余代码无需任何改动，即无论使用 GPU 还是 CPU，代码保持一致。

多 GPU 支持

Darts 利用 Lightning 的多 GPU 能力，以充分利用可扩展硬件。

多 GPU 训练存在多种并行化策略，这些策略因多进程和数据处理方式的差异，与执行环境紧密相关。

目前 Darts 已测试 ddp_spawn 分布式策略。

根据 Lightning 文档的描述，该策略有一些值得注意的限制，例如无法在以下环境中运行：

Jupyter Notebook、Google COLAB、Kaggle 等。
若您拥有无根包（root package）的嵌套脚本

这意味着，实际执行必须在一个独立的 .py 脚本中进行，且脚本需在训练代码周围包含以下通用上下文：

import torch

if __name__ == '__main__':

    torch.multiprocessing.freeze_support()

即使您的执行环境并非 Windows，main 模式也是必需的（参见此链接）。

除此之外，无需对模型进行其他重大修改，只需在 pl_trainer_args 中启用多 GPU 训练即可，例如：

pl_trainer_kwargs = {"accelerator": "gpu", "devices": -1, "auto_select_gpus": True}

该方法会自动选择所有可用 GPU 进行训练。也可以手动设置设备数量。

ddp 系列策略会为每个 GPU 创建独立子进程，因此内存内容（特别是 Dataloader）会被复制。因此，根据 Lightning 文档的描述，建议谨慎设置 Dataloader(num_workers=N) 的 N 值过大：

“Dataloader(num_workers=N)，其中 N 较大，会使 DDP 训练成为瓶颈……即：会非常缓慢甚至完全无法运行。这是 PyTorch 的限制。”

其他分布式策略在 Darts 中的使用目前可能可行，但尚未测试，需视具体设置和实验情况而定。

使用 TPU

张量处理单元（TPU）是 Google 专为神经网络机器学习开发的专用集成电路（ASIC）AI 加速器。

获取 TPU 主要有三种方式：

Google Colab
Google Cloud（GCP）
Kaggle

如果您在 Google Colab 这类笔记本中使用 TPU，则应先安装以下依赖：

!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl
!pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchtext==0.10.0 -f https://download.pytorch.org/whl/cu111/torch_stable.html
!pip install pyyaml==5.4.1

然后指示模型使用一个或多个 TPU。以下示例中我们使用 4 个 TPU：

my_model = RNNModel(
    model="RNN",
    ...
    force_reset=True,
    pl_trainer_kwargs={
      "accelerator": "tpu",
      "tpu_cores": [4]
    },
)

输出如下：

WARNING:root:TPU has started up successfully with version pytorch-1.9
GPU available: False, used: False
TPU available: True, using: [4] TPU cores
IPU available: False, using: 0 IPUs

  | Name      | Type    | Params
--------------------------------------
0 | criterion | MSELoss | 0
1 | rnn       | RNN     | 460
2 | V         | Linear  | 21
--------------------------------------
481       Trainable params
0         Non-trainable params
481       Total params
0.002     Total estimated model params size (MB)
Epoch 299: 100% 8/8 [00:00<00:00, 8.52it/s, loss=0.00285, v_num=logs]
<darts.models.forecasting.rnn_model.RNNModel at 0x7ff1b5e4d4d0>

由输出可见，我们的模型正在使用 4 个 TPU。