Google Colab是一个非常方便的Jupyter学习环境,但因为一些限制,每次都需要重新配置环境、上传文件,而且免费版最长只能运行12个小时,所以需要自己保存checkpoint以便下次恢复。
配置环境就不说了,写个脚本每次都跑一下就完了。
上传文件是个问题,如果有几十上百G的文件,总不能每次都重新上传吧?那可老费时间了,好在Google提供了API,可以挂载Google Drive,这样我们只需要上传一次文件到Google Drive就好了。
在使用Google Colab时,可以通过如下步骤挂载自己的云端硬盘:
- 在Google Colab中,点击左侧的Files按钮。
- 在弹出的对话框中,选择Mount Drive。此操作会在当前的notebook中插入一段Python代码
from google.colab import drive
drive.mount('/content/drive')
- 运行这段代码
- 按提示授权Google Colab访问您的云硬盘。
- 成功挂载后,您的云硬盘会在您的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']