小编ali*_*327的帖子

我目前正在尝试使用 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 …