PyTorch 2.0 公布也有一段时间了，大家是不是曾经开始用了呢？PyTorch 2.0 通过引入 torch.compile，能够显着进步训练和推理速度。与 eagerly 模式相同，编译 API 将模型转换为两头计算图（FX graph），而后以某种形式将其编译为低级计算内核，这样能够进步运行速度。
对于 PyTorch 2.0 而言，你看到的可能是：
“只是用 torch.compile 调用包装它们就能够进步运行速度”
然而其实有许多因素会烦扰计算图编译和 / 或达到所需的性能改良。所以须要调整模型和达到最佳性能可能须要从新设计我的项目或批改一些编码习惯。
在本文中，咱们将演示这个新性能的应用，以及介绍在应用它时可能遇到的一些问题。咱们将分享在调整 torch.compile API 时遇到的问题的几个例子。这些例子并不全面，再理论使用是很可能会遇到此处未提及的问题，并且还要 torch.compile 仍在踊跃开发中，还有改良的空间。
Torch 编译背地有许多翻新技术，包含 TorchDynamo、FX Graph、TorchInductor、Triton 等。咱们不会在这篇文章中深入探讨不同的组件，如果你对这些感兴趣，能够查看 PyTorch 文档，外面介绍的十分具体。
TensorFlow 与 PyTorch 的两个不重要的比照
1、在过来，PyTorch 和 TensorFlow 之间有着显著的区别。PyTorch 应用了 eager execution 模式，TensorFlow 应用了 graph 模式，大家都在各自倒退。但起初 TensorFlow 2 引入了 eager execution 作为默认执行模式，TensorFlow 变得有点像 PyTorch。当初 PyTorch 也引入了本人的 graph 模式解决方案，变得有点像 TensorFlow。TensorFlow 与 PyTorch 的竞争仍在持续，但两者之间的差别正在缓缓隐没。
2、人工智能开发是一个时尚的行业。然而风行的 AI 模型、模型架构、学习算法、训练框架等随工夫变动倒退的。就论文而言，几年前咱们解决的大部分模型都是用 TensorFlow 编写的。然而人们常常埋怨高级 model.fit API 限度了他们的开发灵活性，并且 graph 模式使他们无奈调试。而后就有好多人转向了 PyTorch，他们说，“PyTorch 能够以任何想要的形式构建模型并轻松调试”。然而更灵便的的自定义操作会导致开发的复杂性，PyTorch Lightening 等高级的 API 的呈现就是复制了 model.fit API 的个性，而后同样的人又说还有人说“咱们必须适应 PyTorch Lightening，咱们必须用 torch.compile 减速咱们的训练”。既要灵便，又要简略是不可能同时实现的。
注释开始
上面开始介绍对于如何应用 PyTorch 2 编译 API 的技巧汇合，以及一些你可能面临的潜在问题。使模型适应 PyTorch 的 graph 模式可能须要付出不小的致力。心愿这篇文章能帮忙你更好地评估这一致力，并决定采取这一步的最佳形式。
装置 PyTorch2
从 PyTorch 装置文档来看，装置 PyTorch 2 仿佛与装置任何其余 PyTorch 版本没有什么不同，然而在实践中，可能会遇到一些问题。首先，PyTorch 2.0(截至本文时)须要 Python 3.8 或更高版本。而后就是 PyTorch 2 蕴含以前版本中不存在的包依赖项(最显著的是 PyTorch-triton，这是什么我也不晓得，哈)，须要留神可能会会引入新的抵触。
所以如果你对 Docker 相熟，倡议间接应用容器，这样会简略很多。
PyTorch2 兼容性
PyTorch2 的长处之一是它齐全向后兼容，所以咱们即便不应用 torch.compile，依然能够应用 PyTorch 2.0 并从其余新性能和加强中受害。最多就是享受不到速度的晋升，然而不会有兼容性的问题。然而如果你想进一步晋升速度，那么请往下看。
简略例子
让咱们从一个简略的图像分类模型的例子开始。在上面的代码块中，咱们应用 timm Python 包 (版本 0.6.12) 构建一个根本的 Vision Transformer (ViT)模型，并在一个假数据集上训练它 500 步（不是轮次）。这里定义了 use_compile 标记来管制是否执行模型编译 (torch.compile)，use_amp 来管制是应用主动混合精度(AMP) 还是全精度 (FP) 运行。
importtime, os
importtorch
fromtorch.utils.dataimportDataset
fromtimm.models.vision_transformerimportVisionTransformer
use_amp=True# toggle to enable/disable amp
use_compile=True# toggle to use eager/graph execution mode
# use a fake dataset (random data)
classFakeDataset(Dataset):
def__len__(self):
return1000000
def__getitem__(self, index):
rand_image=torch.randn([3, 224, 224], dtype=torch.float32)
label=torch.tensor(data=[index%1000], dtype=torch.int64)
returnrand_image, label
deftrain():
device=torch.cuda.current_device()
dataset=FakeDataset()
batch_size=64
# define an image classification model with a ViT backbone
model=VisionTransformer()
ifuse_compile:
model=torch.compile(model)
model.to(device)
optimizer=torch.optim.Adam(model.parameters())
data_loader=torch.utils.data.DataLoader(dataset,
batch_size=batch_size, num_workers=4)
loss_function=torch.nn.CrossEntropyLoss()
t0=time.perf_counter()
summ=0
count=0
foridx, (inputs, target) inenumerate(data_loader, start=1):
inputs=inputs.to(device)
targets=torch.squeeze(target.to(device), -1)
optimizer.zero_grad()
withtorch.cuda.amp.autocast(
enabled=use_amp,
dtype=torch.bfloat16
):
outputs=model(inputs)
loss=loss_function(outputs, targets)
loss.backward()
optimizer.step()
batch_time=time.perf_counter() -t0
ifidx>10: # skip first few steps
summ+=batch_time
count+=1
t0=time.perf_counter()
ifidx>500:
break
print(f'average step time: {summ/count}')
if__name__=='__main__':
train()在下表记录了比拟性能后果。这些后果依据环境不同而有很大的变动，所以及供参考
能够看到，应用 AMP(28.6%)比应用 FP(4.5%)时，模型编译带来的性能晋升要显著得多。这是一个家喻户晓的差别。如果你还没有应用 AMP 进行训练，那么其实对于训练速度的晋升是从 FP 过渡到 AMP，所以先举荐你应用 AMP。另外就是性能晋升随同着 GPU 内存利用率的十分轻微的减少。
当扩大到多个 gpu 时，因为在编译图上实现分布式训练的形式，比拟性能可能会发生变化。具体细节看官网文档。
https://pytorch.org/get-started/pytorch-2.0/#distributed
高级选项
compile API 蕴含许多用于管制 graph 创立的选项，可能针对特定模型对编译进行微调，并可能进一步提高性能。上面的代码块是官网的函数介绍：
defcompile(model: Optional[Callable] =None, *,
fullgraph: builtins.bool=False,
dynamic: builtins.bool=False,
backend: Union[str, Callable] ="inductor",
mode: Union[str, None] =None,
options: Optional[Dict[str, Union[str, builtins.int, builtins.bool]]] =None,
disable: builtins.bool=False) ->Callable:
"""
Optimizes given model/function using TorchDynamo and specified backend.
Args:
model (Callable): Module/function to optimize
fullgraph (bool): Whether it is ok to break model into several subgraphs
dynamic (bool): Use dynamic shape tracing
backend (str or Callable): backend to be used
mode (str): Can be either "default", "reduce-overhead" or "max-autotune"
options (dict): A dictionary of options to pass to the backend.
disable (bool): Turn torch.compile() into a no-op for testing
"""mode 编译模式: 容许您在最小化编译所需的开销 (“reduce-overhead”) 和最大化潜在的性能晋升 (“max-autotune”) 之间进行抉择。
下表比拟了不同编译模式下编译上述 ViT 模型的后果。
能够看到编译模式的行为与命名的十分类似，“reduce-overhead”以额定的内存利用为代价缩小了编译工夫，“max-autotune”以高编译工夫开销为代价取得了最佳性能。
backend 编译器后端:API 应用哪个后端将两头示意 (IR) 计算图 (FX graph) 转换为低级内核操作。这个选项对于调试 graph 编译问题和更好地了解 torch.compile 的外部十分有用。在大多数状况下，默认的 Inductor 后端仿佛可能提供最佳的训练性能后果。有很多后端列表，咱们能够应用上面命令查看：
fromtorchimport_dynamo
print(_dynamo.list_backends())咱们测试应用 nvprims-nvfuser 后端，能够取得比 eager 模式 13% 的性能晋升(与默认后端 28.6% 的性能晋升相比)。具体区别还是要看 Pytorch 文档，咱们这里就不细说了，因为文档都有。
fullgraph 强制单个图: 这个参数是十分有用，能够确保没有任何不心愿的图截断。
dynamic 动静形态: 目前 2.0 对具备动静形态的张量的编译反对在某种程度上是无限的。编译具备动静形态的模型的一个常见解决方案是从新编译，但会大大增加开销并大大降低训练速度。如果您的模型的确蕴含动静形态，将动静标记设置为 True 将带来更好的性能，特地是缩小从新编译的次数。
都有什么是动静形态呢，最简略的就是工夫序列或文本长度不同，如果不进行对齐操作的话序列长度不同就是动静的形态。
性能剖析
PyTorch Profiler 是用来剖析 PyTorch 模型性能的要害工具之一，能够评估和剖析图编译优化训练步骤的形式。在上面的代码块中，咱们用 profiler 生成 TensorBoard 的后果，来查看训练的性能：
out_path=os.path.join(os.environ.get('SM_MODEL_DIR','/tmp'),'profile')
fromtorch.profilerimportprofile, ProfilerActivity
withprofile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(
wait=20,
warmup=5,
active=10,
repeat=1),
on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name=out_path)
) asp:
foridx, (inputs, target) inenumerate(data_loader, start=1):
inputs=inputs.to(device)
targets=torch.squeeze(target.to(device), -1)
optimizer.zero_grad()
withtorch.cuda.amp.autocast(
enabled=use_amp,
dtype=torch.bfloat16
):
outputs=model(inputs)
loss=loss_function(outputs, targets)
loss.backward()
optimizer.step()
p.step()下图是从 PyTorch Profiler 生成的 TensorBoard 中截取的。它提供了在下面编译模型试验的训练步骤中在 GPU 上运行的内核的详细信息。
咱们可能看到 torch.compile 减少了 GPU 张量外围的利用率(从 51% 到 60%)，并且它引入了应用 Triton 开发的 GPU 内核。
调试模型编译问题
torch.compile 目前处于测试阶段，如果你遇到问题，并且侥幸的话，会失去一个信息谬误，咱们能够间接搜寻解决，或者问问 chatgpt。然而如果你不那么侥幸，就须要本人寻找问题的本源。
这里解决编译问题的次要资源是 TorchDynamo 故障排除文档，其中包含调试工具列表并提供诊断谬误的分步指南。然而目前这些工具和技术仿佛更多地针对 PyTorch 开发人员而不是 PyTorch 用户的。它们兴许能够帮忙解决导致编译问题的基本问题，然而十分大的可能是它们实际上跟本没有任何帮忙，那怎么办呢？
这里咱们演示一个自行解决问题的过程，依照这样的思路，能够解决一些问题。
上面是一个简略的分布式模型，其中包含对 torch.distributed.all_reduce 的调用。模型在 eager 模式下按预期运行，但在 graph 编译期间失败并呈现“attribute error”（torch.classes.c10d.ProcessGroup does not have a field with name‘shape’）。咱们须要将日志级别进步到 INFO，而后发现发现错误在计算的“第 3 步”中，即 TorchInductor。而后通过验证“eager”和“aot_eager”后端的编译是否胜利，最初创立一个最小的代码示例，应用 PyTorch Minifier 重现失败。
importos, logging
importtorch
fromtorchimport_dynamo
# enable debug prints
torch._dynamo.config.log_level=logging.INFO
torch._dynamo.config.verbose=True
# uncomment to run minifier
# torch._dynamo.config.repro_after="aot"
defbuild_model():
importtorch.nnasnn
importtorch.nn.functionalasF
classDumbNet(nn.Module):
def__init__(self):
super().__init__()
self.conv1=nn.Conv2d(3, 6, 5)
self.pool=nn.MaxPool2d(2, 2)
self.fc1=nn.Linear(1176, 10)
defforward(self, x):
x=self.pool(F.relu(self.conv1(x)))
x=torch.flatten(x, 1)
x=self.fc1(x)
withtorch.no_grad():
sum_vals=torch.sum(x,0)
# this is the problematic line of code
torch.distributed.all_reduce(sum_vals)
# add noise
x=x+0.1*sum_vals
returnx
net=DumbNet()
returnnet
deftrain():
os.environ['MASTER_ADDR'] =os.environ.get('MASTER_ADDR',
'localhost')
os.environ['MASTER_PORT'] =os.environ.get('MASTER_PORT',
str(2222))
torch.distributed.init_process_group('nccl', rank=0,
world_size=1)
torch.cuda.set_device(0)
device=torch.cuda.current_device()
model=build_model()
model=torch.compile(model)
# replace with this to verfiy that error is not in TorchDynamo
# model = torch.compile(model, 'eager')
# replace with this to verfiy that error is not in AOTAutograd
# model = torch.compile(model, 'aot_eager')
model.to(device)
rand_image=torch.randn([4, 3, 32, 32], dtype=torch.float32).to(device)
model(rand_image)
if__name__=='__main__':
train()在这个的示例中，运行生成的 minifier_launcher.py 脚本会导致不同的属性谬误（比方 Repro’object has no attribute‘_tensor_constant0’），这个对于咱们的演示没有太大帮忙，咱们临时疏忽他，这也阐明了，torch.compile 还不欠缺，还须要更大的改良空间，或者说如果解决不要问题，那就别用了，至多“慢”要比不能用好，对吧（而且速度晋升也无限）
常见的图截断问题
Pytorch eager 模式劣势之一是可能将纯 Pythonic 代码与 PyTorch 操作交错在一起。然而这种自在在应用 torch.compile 时受到很大限度。因为 Pythonic 操作导致 TorchDynamo 将计算图拆分为多个组件，从而妨碍了性能晋升的后劲。而咱们代码优化的指标是尽可能减少此类图截断。最简略的方法是用 fullgraph 标记编译模型。这杨能够提醒删除导致图截断的任何代码，而且还会通知咱们如何最好地适应 PyTorch2 的开发习惯。然而要运行分布式代码，则必须将他设为 False，因为以后实现 GPU 之间通信的形式须要图拆分。咱们也能够应用 torch._dynamo.explain 序来剖析图截断。
以下代码块演示了一个简略模型，在其前向传递中有四个潜在的图截断，然而这种在应用形式在典型的 PyTorch 模型中并不少见。
importtorch
fromtorchimport_dynamo
importnumpyasnp
defbuild_model():
importtorch.nnasnn
importtorch.nn.functionalasF
classDumbNet(nn.Module):
def__init__(self):
super().__init__()
self.conv1=nn.Conv2d(3, 6, 5)
self.pool=nn.MaxPool2d(2, 2)
self.fc1=nn.Linear(1176, 10)
self.fc2=nn.Linear(10, 10)
self.fc3=nn.Linear(10, 10)
self.fc4=nn.Linear(10, 10)
self.d= {}
defforward(self, x):
x=self.pool(F.relu(self.conv1(x)))
x=torch.flatten(x, 1)
asserttorch.all(x>=0) # graph break
x=self.fc1(x)
self.d['fc1-out'] =x.sum().item() # graph break
x=self.fc2(x)
forkinnp.arange(1): # graph break
x=self.fc3(x)
print(x) # graph break
x=self.fc4(x)
returnx
net=DumbNet()
returnnet
deftrain():
model=build_model()
rand_image=torch.randn([4, 3, 32, 32], dtype=torch.float32)
explanation=torch._dynamo.explain(model, rand_image)
print(explanation)
if__name__=='__main__':
train()图截断不会导致编译失败(除非设置了 fullgraph 标记)。所以很有可能模型正在编译和运行，但实际上蕴含多个图截断，这会减慢它的速度。
训练问题故障排除
在目前来说，应用 Pytorch2 胜利编译的模型就能够认为是一项值得庆贺的成就，但这并不能保障训练肯定会胜利。
在 GPU 上运行的低级内核在 eager 模式和 graph 模式之间会有所不同。某些高级操作可能会体现出不同的行为。你可能会发现在 eager 模式下运行的操作在 graph 模式下会失败（例如 torch.argmin）。或者会发现计算中的数值差别会影响训练。
graph 模式下的调试比 eager 模式下的调试艰难得多。在 eager 模式下，每一行代码都是独立执行的，咱们能够在代码中的任意点搁置断点取得前张量值。而在 graph 模式下，代码定义的模型在解决之前会经验屡次转换，设置的断点可能不会被触发。
所以能够先应用 eager 模式，模型跑通当前，再将 torch.compile 别离利用于每个局部，或者通过插入打印和 / 或 Tensor.numpy 调用来生成图截断，这样咱们可能会会胜利触发代码中的断点。也就是说如果用 torch.compile 的话对于开发来说，要消耗更长的工夫，所以训练和开发速度的取舍就要看你本人的抉择了。
然而别忘了咱们下面说的你的模型在增加了 torch.compile 后也不肯定能正确运行，这又是一个有形的老本。
在图中蕴含损失函数
通过应用 torch.compile 调用包装 PyTorch 模型 (或函数) 来启用 graph 模式。然而损失函数不是编译调用的一部分，也不是生成图的一部分。所以损失函数是训练步骤中绝对较小的一部分，如果应用 eager 模式运行它不会产生太多开销。然而如果有一个计算量他别大的损失函数，也是能够通过将其蕴含在编译的计算图中来进一步提高性能的。
在上面的代码中，咱们定义了一个损失函数，用于执行从大型 ViT 模型 (具备 24 个 ViT 块) 到较小的 ViT 模型 (具备 12 个 ViT 块) 的模型蒸馏。
importtorch
fromtimm.models.vision_transformerimportVisionTransformer
classExpensiveLoss(torch.nn.Module):
def__init__(self):
super(ExpensiveLoss, self).__init__()
self.expert_model=VisionTransformer(depth=24)
iftorch.cuda.is_available():
self.expert_model.to(torch.cuda.current_device())
self.mse_loss=torch.nn.MSELoss()
defforward(self, input, outputs):
expert_output=self.expert_model(input)
returnself.mse_loss(outputs, expert_output)这是一个比 CrossEntropyLoss 计算量大得多的损失函数，这里又 2 种办法让他执行的更快，
1、loss 函数封装在 torch.compile 调用中，如下所示:
loss_function=ExpensiveLoss()
compiled_loss=torch.compile(loss_function)这个办法的毛病是损失函数的编译图与模型的编译图不相交，然而它的长处非常明显，就是简略。
2、创立一个蕴含模型和损失的包装器模型来将模型和损失一起编译，并将后果损失作为输入返回。
importtime, os
importtorch
fromtorch.utils.dataimportDataset
fromtorchimportnn
fromtimm.models.vision_transformerimportVisionTransformer
# use a fake dataset (random data)
classFakeDataset(Dataset):
def__len__(self):
return1000000
def__getitem__(self, index):
rand_image=torch.randn([3, 224, 224], dtype=torch.float32)
label=torch.tensor(data=[index%1000], dtype=torch.int64)
returnrand_image, label
# create a wrapper model for the ViT model and loss
classSuperModel(torch.nn.Module):
def__init__(self):
super(SuperModel, self).__init__()
self.model=VisionTransformer()
self.expert_model=VisionTransformer(depth=24iftorch.cuda.is_available() else2)
self.mse_loss=torch.nn.MSELoss()
defforward(self, inputs):
outputs=self.model(inputs)
withtorch.no_grad():
expert_output=self.expert_model(inputs)
returnself.mse_loss(outputs, expert_output)
# a loss that simply passes through the model output
classPassthroughLoss(nn.Module):
def__call__(self, model_output):
returnmodel_output
deftrain():
device=torch.cuda.current_device()
dataset=FakeDataset()
batch_size=64
# create and compile the model
model=SuperModel()
model=torch.compile(model)
model.to(device)
optimizer=torch.optim.Adam(model.parameters())
data_loader=torch.utils.data.DataLoader(dataset,
batch_size=batch_size, num_workers=4)
loss_function=PassthroughLoss()
t0=time.perf_counter()
summ=0
count=0
foridx, (inputs, target) inenumerate(data_loader, start=1):
inputs=inputs.to(device)
targets=torch.squeeze(target.to(device), -1)
optimizer.zero_grad()
withtorch.cuda.amp.autocast(
enabled=True,
dtype=torch.bfloat16
):
outputs=model(inputs)
loss=loss_function(outputs)
loss.backward()
optimizer.step()
batch_time=time.perf_counter() -t0
ifidx>10: # skip first few steps
summ+=batch_time
count+=1
t0=time.perf_counter()
ifidx>500:
break
print(f'average step time: {summ/count}')
if__name__=='__main__':
train()这种办法的毛病是，当在推理模式下运行模型时，须要从包装器模型中提取外部的理论模型。
这两种选项的性能晋升幅度大致相同都是 8%，也就是说，对 loss 进行编译也是优化的一个重要局部。
动静形态
官网也说了 torch.compile 对动静形态的模型的编译反对是无限的。compile API 蕴含 dynamic 参数，用于向编译器发出信号，然而这种形式对于性能晋升帮忙的水平是值得狐疑的。如果你正在尝试编译和优化动态图并面临问题，那么还是不要应用 torch.compile，因为太麻烦了。
总结
PyTorch 2.0 编译模式具备显著进步训练和推理速度的后劲，能够显著节省成本，然而模型实现这一后劲所需的工作量可能会有很大差别。许多公共模型只须要批改一行代码。而其余模型特地是那些蕴含非标准操作、动静形态和 / 或大量交织 Python 代码的模型，可能得失相当甚至无奈进行。然而当初开始批改模型是一个很好的抉择，因为目前来看 torch.compile 对于 PyTorch2 来说是一个重要且继续的个性。
https://avoid.overfit.cn/post/dfea563957fc43a19f1aaf7733888031
作者：Chaim Rand