我应该如何正确使用torch.compile？

Question

我目前正在尝试使用 pytorch 2.0 来提高我的项目的训练性能。我听说 torch.compile 可能会增强某些模型。

所以我的问题（目前）很简单；我应该如何将 torch.compile 与大型模型一起使用？

例如，我应该像这样使用 torch.model 吗？

class BigModel(nn.Module):
    def __init__(self, ...):
        super(BigModel, self).__init__()
        self.model = nn.Sequential(
            SmallBlock(), 
            SmallBlock(), 
            SmallBlock(), 
            ...
        )
        ...

class SmallBlock(nn.Module):
    def __init__(self, ...):
        super(SmallBlock, self).__init__()
        self.model = nn.Sequential(
            ...some small model...
        )

model = BigModel()
model_opt = torch.compile(model)

，或者像这样？

class BigModel(nn.Module):
    def __init__(self, ...):
        super(BigModel, self).__init__()
        self.model = nn.Sequential(
            SmallBlock(), 
            SmallBlock(), 
            SmallBlock(), 
            ...
        )
        ...

class SmallBlock(nn.Module):
    def __init__(self, ...):
        super(SmallBlock, self).__init__()
        self.model = nn.Sequential(
            ...some small model...
        )
        self.model = torch.compile(self.model)

model = BigModel()
model_opt = torch.compile(model)

总结一下，

1. 应该编译每一层吗？或者 torch.compile 自动执行此操作？
2. 有什么正确使用 torch.compile 的技巧吗？

说实话，我都试过了，但没有什么区别。

而且，它并没有显着加速，我只是检查了我的模型的加速率约为 5 ~ 10%。

Answer 1

torch.compile 的默认模式似乎不起作用，但它有另一种模式可以真正加速您的模型。""" torch.compile(${yourmodel}, mode="reduce-overhead") """