Могут ли модели, обученные PyTorch, передаваться между графическими процессорами и ТПУ?

#pytorch #gpu #tpu

Вопрос:

После использования графического процессора для некоторых для обучения модели PyTorch, могу ли я использовать сохраненные веса для продолжения обучения моей модели на TPU?

Комментарии:

1. У Pytorch нет поддержки TPU

2. @talonmies github.com/pytorch/xla

Ответ №1:

После использования GPU в течение некоторого времени могу ли я использовать сохраненные веса для обучения своей модели с использованием TPU?

Да, если вы сохранили свою обученную графическому процессору модель с помощью, скажем

     torch.save(model.save_dict(), 'model.pt')
 

вы можете загрузить его снова для использования на ТПУ (используя https://github.com/pytorch/xla) в отдельной программе, выполняемой с

     import torch_xla.utils.serialization as xser
    
    model.load_state_dict(xser.load('model.pt'))