CPU加速

1.调用一个 JIT 解释器

PyTorch 允许在 Torchscript 模型推理期间使用多个 CPU 线程。下图显示了在典型应用程序中可以找到的不同级别的并行性：

一个或多个推理线程在给定输入上执行模型的前向传递。每个推理线程调用一个 JIT 解释器，该解释器一个接一个地执行模型内联的操作。

2.使用ATen、MKL 和 MKL-DNN来构建

PyTorch 使用内部 ATen 库来实现操作。除此之外，PyTorch 还可以构建支持外部库，例如MKL和MKL-DNN，以加快 CPU 上的计算。

ATen、MKL 和 MKL-DNN 支持操作内并行性，并依赖OpenMP和TBB并行化库来实现它 

这有两篇关于MKL-DNN加速  使用mkldnn为cpu运行网络加速_绝望的乐园-CSDN博客_mkldnn

多卡或多机加速

torch支持多进程的分布数据的并行运算  DDP  torch.nn.parallel.DistributedDataParallel

内部实现原理可参照Distributed Data Parallel — PyTorch 1.10 documentation

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # forward pass
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)
    # backward pass
    loss_fn(outputs, labels).backward()
    # update parameters
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    main()