2024年3月18日17:49:10评论3 views字数 7635阅读25分27秒阅读模式

技术分享：torch.compile在GammaGL上的实践

本次技术分享主要聚焦于 PyTorch 2.0 版本更新的内容在 GammaGL 上的应用，带领读者进一步了解 torch.compile 的使用方法以及在 GammaGL 上的应用实践。

1 引言

目前，研究人员选择的主流深度学习编程框架为 PyTorch，PyTorch 为研究人员提供了以下优秀的功能：

简单、高效的编程接口
自动微分
数据处理与管理
模型部署与训练
硬件加速

随着开发人员的优化推进，PyTorch 在 2.0版本推出最新的优化方法：torch.compile，它采用 JIT (Just In Time) 编译的方式将 PyTorch 代码转化为高效的内核，而用户无需修改很多的代码。利用该方法，用户可以构建更加高效的代码，缩短模型训练的运行时间，有助于研究人员的实验验证。

专用于优化图神经网络实现的图神经网络框架大多建立在深度学习框架之上，例如 PyTorch-Geometric (PyG) 建立在PyTorch之上，Deep Graph Library (DGL) 建立在 PyTorch, TensorFlow, MXNet 之上。由 GAMMA Lab 建设和维护的支持多后端开源图神经网络框架 GammaGL 在 0.4版本以后全面兼容 PyTorch 2.0。如何充分有效地利用深度学习框架带来的优势，逐渐成为研究人员和开发人员面临的挑战。

因此，本次技术分享首先将介绍 PyTorch 2.0 版本中torch.compile的功能和使用方法，并尝试在 GammaGL 的训练中引入该方法提高图神经网络模型训练的效率。

2 torch.compile

PyTorch 在 2.0 版本更新了torch.compile接口，相比于之前的 TorchScript 和 FX Tracing 具有更多的优势。

安装依赖如下：

torch >= 2.0
CUDA >= 11.7
numpy
scipy

接口介绍：

torch.compile(model=None, *, fullgraph=False, dynamic=None, backend='inductor', mode=None, options=None, disable=False)

参数

model (Callable): 要优化的模块或函数
fullgraph (bool): 默认为False，则 torch.compile 会尝试发现可优化区域。如果设置为True，则要求用户必须传递单一的模型计算图，否则报错。
dynamic (bool or None): 默认为None，则 torch.compile 会自动判断是否编译为动态内核。如果设置为True，则会采用动态编译，目标是针对不同输入尺寸的数据无需多次编译，有效地适应各种尺寸数据，充分利用计算资源。如果设置为False，则采用静态编译，固定模型输入数据的尺寸，达到最优性能；一旦改变输入数据尺寸，则会再次编译。
backend (str or Callable): PyTorch 使用的底层后端，默认为"inductor"，含义为性能和开销之间进行平衡。非实验性（稳定版本）后端可以通过torch._dynamo.list_backends()查看，实验性（非稳定版本）后端可以通过torch._dynamo.list_backends(None)查看。用户自定义算法可以参照：custome_backend。
mode (str): 可以是 "default"、"reduce-overhead"、"max-autotune" 或 "max-autotune-no-cudagraphs" 中的任意一个。

"default" 是默认模式，平衡了性能和开销。
"reduce-overhead" 模式减少了小批量时 Python 的开销，可能会因为缓存工作空间内存而使用更多内存。
"max-autotune" 模式利用基于 Triton 的矩阵乘法和卷积，默认启用 CUDA 图，一般需要花费更长的编译时间获得更好的性能。
"max-autotune-no-cudagraphs" 类似于 "max-autotune"，但不使用 CUDA 图。

options (dict): 传递给后端的选项字典。
disable (bool): 默认为False，如果设置为True，则会取消使用编译功能。

基本使用方法：

import torch

def foo(x, y):
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b
opt_foo1 = torch.compile(foo)
print(opt_foo1(torch.randn(10, 10), torch.randn(10, 10)))

或

import torch

@torch.compile
def opt_foo2(x, y):
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b
print(opt_foo2(torch.randn(10, 10), torch.randn(10, 10)))

速度测试：

测试环境：

CPU：Intel(R) Xeon(R) Gold 6348 CPU @ 2.60GHz
GPU：NVIDIA GeForce RTX 3090 (24GB)

首先，我们使用PyTorch定义一个卷积神经网络：

import torch
import torch.nn as nn
import torch.nn.functional as F

class LeNet5(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, padding=2)
        self.pool1 = nn.MaxPool2d((2, 2))
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool2 = nn.MaxPool2d((2, 2))
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(len(x), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

紧接着，我们定义模型并测试：

import time

x = torch.randn(64, 1, 28, 28).cuda()
label = torch.randint(0, 10, (64, )).cuda()

model = LeNet5().cuda()
model_opt = torch.compile(model, mode="reduce-overhead")
start = time.time()
model(x)
end = time.time()
print("eager:", end - start)
start = time.time()
model_opt(x)
end = time.time()
print("compile:", end - start)

测试结果如下：

eager: 0.16981077194213867
/home/lyq/anaconda3/envs/turtle/lib/python3.10/site-packages/torch/_inductor/compile_fx.py:140: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
compile: 2.440474510192871

编译消耗的时间更长，这很正常，因为我们仅仅运行了一次模型，在深度学习操作中，通常需要进行多次迭代，接下来我们看一下在一个10次循环中的时间消耗情况：

import time

EPOCH = 10

x = torch.randn(256, 1, 28, 28).cuda()
label = torch.randint(0, 10, (64, )).cuda()

model = LeNet5().cuda()
import torch._dynamo
torch._dynamo.reset()
model_opt = torch.compile(model, mode="reduce-overhead")

eager_times = []
for _ in range(EPOCH):
    start = time.time()
    model(x)
    end = time.time()
    print("eager eval time:", end - start)
    eager_times.append(end - start)

compile_times = []
for _ in range(EPOCH):
    start = time.time()
    model_opt(x)
    end = time.time()
    print("compile eval time:", end - start)
    compile_times.append(end - start)

import numpy as np
eager_med = np.median(eager_times)
compile_med = np.median(compile_times)
speedup = eager_med / compile_med
print(f"(eval) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x")

测试结果如下：

eager eval time: 0.1669483184814453
eager eval time: 0.0004298686981201172
eager eval time: 0.00034618377685546875
eager eval time: 0.0003509521484375
eager eval time: 0.0003256797790527344
eager eval time: 0.0003223419189453125
eager eval time: 0.0003314018249511719
eager eval time: 0.00032520294189453125
eager eval time: 0.0003325939178466797
eager eval time: 0.0003235340118408203
/home/lyq/anaconda3/envs/turtle/lib/python3.10/site-packages/torch/_inductor/compile_fx.py:140: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
compile eval time: 2.394608974456787
compile eval time: 0.06615447998046875
compile eval time: 0.00041675567626953125
compile eval time: 0.0002601146697998047
compile eval time: 0.0002110004425048828
compile eval time: 0.00020885467529296875
compile eval time: 0.00019621849060058594
compile eval time: 0.00019097328186035156
compile eval time: 0.00019168853759765625
compile eval time: 0.00019168853759765625
(eval) eager median: 0.0003319978713989258, compile median: 0.00020992755889892578, speedup: 1.581487791027825x

在该样例中，编译后的代码，可以有更短的时间消耗。不过，是否选择合适的mode，会在很大程度上影响到编译后代码的运行效率。

更多详细的使用方法以及注意问题，可参照 Introduction to torch.compile.

3 GammaGL实践

接下来，我们使用 GammaGL 的 PyTorch 后端实现 GAT 模型，配合torch.compile给出一个完整的图神经网络训练流程的实验结果，并给出速度测试。参考代码可详见于：https://github.com/dddg617/GammaGL/tree/gat/examples/compile。

安装流程：

$ pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu118
$ pip install git+https://github.com/dddg617/tensorlayerx.git@nightly
$ pip install gammagl-pt21==0.4.0.post118

模型定义及编译：

import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['TL_BACKEND'] = 'torch'
from gammagl.datasets import Planetoid
from gammagl.layers.conv import GATConv
from gammagl.utils import add_self_loops, mask_to_index
import torch
import time
import tensorlayerx as tlx
import torch._dynamo
torch._dynamo.reset()

class GATModel(tlx.nn.Module):

    def __init__(self,feature_dim,
                 hidden_dim,
                 num_class,
                 heads,
                 drop_rate,
                 name=None,
                 ):
        super().__init__(name=name)
        
        self.conv1 = GATConv(
            in_channels=feature_dim,
            out_channels=hidden_dim,
            heads=heads,
            dropout_rate=drop_rate,
            concat=True
        )
        self.conv2 = GATConv(
            in_channels=hidden_dim * heads,
            out_channels=num_class,
            heads=heads,
            dropout_rate=drop_rate,
            concat=False
        )
        self.elu = tlx.layers.ELU()
        self.dropout = tlx.layers.Dropout(drop_rate)

    def forward(self, x, edge_index, num_nodes):
        x = self.dropout(x)
        x = self.conv1(x, edge_index, num_nodes)
        x = self.elu(x)
        x = self.dropout(x)
        x = self.conv2(x, edge_index, num_nodes)
        x = self.elu(x)
        return x

model = GATModel(
            feature_dim=dataset.num_node_features,
            hidden_dim=args.hidden_dim,
            num_class=dataset.num_classes,
            heads = args.heads,
            drop_rate=args.drop_rate,
            name="GAT"
            )
model = torch.compile(model, mode="reduce-overhead", dynamic=False)
optimizer = torch.optim.Adam(params=parameters(model), lr=args.lr, weight_decay=args.l2_coef)

实验设置：

hidden_dim: 256
num_epoch: 1000
heads: 32
dataset: pubmed

实验结果：

eager mode time: 196.90568041801453
compile mode time: 191.0114598274231
eager mode each epoch time: near 0.19
compile mode each epoch time: near 0.17

通常，在模型复杂度越高、训练迭代次数越多的训练流程、输入数据量越多的场景中，使用torch.compile可以带来更大的收益，一般类似于GCN, GAT等比较简单的图神经网络模型，并不推荐。随着目前深度学习模型越来越复杂，预先尝试torch.compile并调出合适的编译设置，往往可以带来“磨刀不误砍柴工”的效果。

实验室算法库简介

🎉两算法库均获得OpenI启智社区“社区优秀孵化项目奖”

GammaGL：GammaGL是一个基于TensorLayerX实现的开源的多框架的图神经网络（GNN）算法平台，该框架抽象出图数据、消息传递、采样等组件，已经实现了二十余种图数据集和四十余种术界经典GNN算法。

GammaGL发布0.4版本
https://github.com/BUPT-GAMMA/GammaGL

OpenHGNN：OpenHGNN是一个基于深度学习框架PyTorch和图神经网络框架DGL作为底层框架的异质图神经网络（HGNN）工具包。该工具包专注于HGNN的算法模型，抽象出了数据集管理、算法流程、任务评测等模块，提供了超参数调优组件和最优的超参数设置，发布了HGNN设计空间和Benchmark等。

OpenHGNN发布0.6版本
OpenHGNN：https://github.com/BUPT-GAMMA/OpenHGNN

欢迎 Star ❤ ~

原文始发于微信公众号（北邮 GAMMA Lab）：技术分享：torch.compile在GammaGL上的实践

左青龙
微信扫一扫

右白虎
微信扫一扫

技术分享：torch.compile在GammaGL上的实践

技术分享：torch.compile在GammaGL上的实践

1 引言

2 torch.compile

接口介绍：

基本使用方法：

速度测试：

3 GammaGL实践

安装流程：

模型定义及编译：

实验设置：

实验结果：

实验室算法库简介

🎉两算法库均获得OpenI启智社区“社区优秀孵化项目奖”

GammaGL：GammaGL是一个基于TensorLayerX实现的开源的多框架的图神经网络（GNN）算法平台，该框架抽象出图数据、消息传递、采样等组件，已经实现了二十余种图数据集和四十余种术界经典GNN算法。

欢迎 Star ❤ ~

记一次有源码的渗透测试

经典：“红日三”服务器群靶机

Google如何利用生成式AI加速事件响应流程

Hack The Box-start point-Archetype笔记

使用 Burp 枚举 REST API

fastjson-BCEL不出网打法原理分析

改变content-type类型所触发的csrf

【OSCP】vivifytech

从零开始的Kubernetes攻防

一次js混淆测试记录

发表评论

在线咨询

微信