使用带有GPU和TPU的Torch模型¶

本节是为 Darts 0.17.0 及更高版本编写的。

我们假设您已经了解Darts中的Torch预测模型。如果您是新手，我们建议您先阅读`Torch预测模型指南 <https://unit8co.github.io/darts/userguide/torch_forecasting_models.html>`_。该指南还包含关于性能建议的部分，我们建议您先阅读。最后，这里还有一个`循环神经网络（RNN）模型示例 <https://unit8co.github.io/darts/examples/04-RNN-examples.html>`_，本节将基于此示例进行讲解。

使用 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中使用加速器，我们可以享受到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环境下发生。

除此之外，除了在 pl_trainer_args 中允许多GPU训练外，您的模型不需要进行其他重大修改。

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

此方法自动选择所有可用的GPU进行训练。手动设置设备数量也是可能的。

ddp 策略系列为每个 GPU 创建单独的子进程，因此内存内容（特别是 Dataloader）会被复制。因此，根据 lightning 文档的描述，在设置 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的笔记本中使用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 预测模型

协变量