使用 GPU 和 TPU 使用 Torch 模型¶

本节内容适用于 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

现在我们像这样读取并缩放数据

# Read data:
series = AirPassengersDataset().load()
series = series.astype(np.float32)

# Create training and validation sets:
train, val = series.split_after(pd.Timestamp("19590101"))

# Normalize the time series (note: we avoid fitting the transformer on the validation set)
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 参数指示我们的模型使用我们机器的 GPU，这些参数表示为 pl_trainer_kwargs 字典，如下所示：

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 等。
如果您有一个没有根包的嵌套脚本

实际上这意味着执行必须在一个单独的 .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 创建单独的子进程，内存内容（特别是 Dataloder）会被复制。因此，根据 lightning docs 的描述，不建议将 Dataloader(num_workers=N) 设置得过高，因为它表示

“Dataloader(num_workers=N)，其中 N 较大，会瓶颈化 DDP 训练……即：会非常慢或根本无法工作。这是 PyTorch 的限制。”

目前在 Darts 中使用其他分发策略可能可行，但尚未经过测试，并且取决于个人设置/实验。

使用 TPU¶

谷歌专门为神经网络机器学习开发了张量处理单元 (TPU)，这是一种 AI 加速器专用集成电路 (ASIC)。

有三种主要方法可以访问 TPU

Google Colab
Google Cloud (GCP)
Kaggle

如果您在 Google Colab 类型的 notebook 中使用 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。在我们的示例中，我们使用四个 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。

Torch 预测模型

协变量