Google Colab挂载云端硬盘并保存checkpoint

Google Colab是一个非常方便的Jupyter学习环境,但因为一些限制,每次都需要重新配置环境、上传文件,而且免费版最长只能运行12个小时,所以需要自己保存checkpoint以便下次恢复。

配置环境就不说了,写个脚本每次都跑一下就完了。

上传文件是个问题,如果有几十上百G的文件,总不能每次都重新上传吧?那可老费时间了,好在Google提供了API,可以挂载Google Drive,这样我们只需要上传一次文件到Google Drive就好了。

在使用Google Colab时,可以通过如下步骤挂载自己的云端硬盘:

  1. 在Google Colab中,点击左侧的Files按钮。
  2. 在弹出的对话框中,选择Mount Drive。此操作会在当前的notebook中插入一段Python代码
from google.colab import drive
drive.mount('/content/drive')
  1. 运行这段代码
  2. 按提示授权Google Colab访问您的云硬盘。
  3. 成功挂载后,您的云硬盘会在您的Colab环境中映射为一个文件夹。

接下来,您可以使用基本的Python文件操作函数(如open、read等)访问您的云盘文件。

例如,您可以使用如下代码读取您云盘中的文件:

with open('/content/drive/My Drive/path/to/your/file', 'r') as f:
    data = f.read()

在Google Colab中保存checkpoint可以使用如下代码:

# your code here
# ...

# save checkpoint
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, '/content/drive/My Drive/path/to/your/checkpoint')

这段代码将您的模型相关信息、优化器相关信息和当前loss保存到一个字典中,并使用PyTorch内置函数torch.save将字典保存到您的云盘中。

在需要恢复模型时,您可以使用如下代码加载checkpoint:

# load checkpoint
checkpoint = torch.load('/content/drive/My Drive/path/to/your/checkpoint')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

Leave a Comment

豫ICP备19001387号-1