Pytorch 在Colab TPU上运行Pyrotch堆叠模型

我正在尝试在Colab多核TPU上运行我的模型，但我真的不知道怎么做。我试过了，但我遇到了一些错误，我无法修复它，但我认为可能还有更简单的等待

关于我的模型：

class BERTModel(nn.Module):
    def __init__(self,...):
        super().__init__()
        if ...:
            self.bert_model = XLMRobertaModel.from_pretrained(...)   # huggingface XLM-R
        elif ...:
            self.bert_model = others_model.from_pretrained(...)   # huggingface XLM-R
        
        ... # some other model's parameters
        
    def forward(self,...):
        bert_input = ...
        output = self.bert_model(bert_input)
        
        ... # some function that process on output
        
    def other_function(self,...):
        # just doing some process on output. like concat layers's embedding and return ...
        
class MAINModel(nn.Module):
    def __init__(self,...):
        super().__init__()
        
        print('Using model 1')
        self.bert_model_1 = BERTModel(...)
        
        print('Using model 2')
        self.bert_model_2 = BERTModel(...)
        
        self.linear = nn.Linear(...)
        
    def forward(self,...):
        bert_input = ...
        bert_output = self.bert_model(bert_input)
        linear_output = self.linear(bert_output)
   
        return linear_output

---

你能告诉我如何在Colab TPU上运行像我的模型这样的模型吗？我使用Colab PRO来确保Ram内存不是一个大问题。非常感谢您。

我将在这里列出以下示例：

也许可以从这样一个简单的模型开始：

在您共享的伪代码中，没有对

torch\u xla

库的引用，这是在TPU上使用PyTorch所必需的。我建议从我共享的目录中的一个工作的Colab笔记本开始，然后用您自己的模型替换部分模型。如果您想在TPU上运行模型，则需要修改在GPU上使用本机PyTorch运行的模型的总体培训代码中的一些位置（通常为3-4）。有关某些更改的说明，请参阅。另一个大的变化是用一个并行加载器包装默认的数据加载器，如示例MNIST colab I shared所示

---

如果您在其中一个COLAB中看到任何特定错误，请随时打开一个问题：

当您共享收到的错误消息时，它总是很有用的。请将完整的stacktrace添加到您的问题中。