way to train&inference with TPU

[doheon114]

Hello. Is there any way to accelerate processes by TPU? (using XLA or Jax)
Thank you for your reply.

[zhangyuxuann]

@doheon114 Protenix do not support training or inference with TPU yet.

[doheon114]

I'm sorry for asking so many questions, but would it be possible to save the checkpoint and structure of this open-source model using torch.save and then load it to train with my own code? It seems like the code might be too intertwined, making it difficult, but since I'm a beginner in coding, I’d like to try if possible.

I want to use an approach like this:

model = torch.load('/home/hu2/model_structure.pth')
  class PretrainedModel(nn.Module):

      def __init__(self, pretrained_model):
          super(PretrainedModel, self).__init__()
          self.pretrained_model = pretrained_model
  
      def forward(self, input_feature_dict, label_full_dict, label_dict):
          outputs = self.pretrained_model(input_feature_dict, label_full_dict, label_dict)
          return outputs
  
  pretrained_model = PretrainedModel(pretrained_model=model)

Would this approach work? Thankyou

[zhangyuxuann]

torch.save(model) saves the entire model structure and its parameters (weights). This is problematic because:
The saved file contains the exact Python class structure, including the file paths where the original class was defined. If your code doesn't have access to the exact same class definition (same file paths, same dependencies), loading will fail.