如何实现Google Colab自动重启并运行带特殊标记的初始化单元格？

核心思路 #

通过脚本识别笔记本中带特殊标记的单元格，将代码保存到临时文件，重启运行时后自动执行该文件，模拟Mathematica的"评估初始化单元格"功能。

具体实现步骤 #

1. 编写初始化单元格识别与重启脚本 #

创建一个代码单元格，用来筛选带#--init--标记的单元格并触发重启：

import json
from google.colab import notebook, runtime

# 读取当前笔记本的全部内容
notebook_json = json.loads(notebook.get_contents())

# 筛选所有带#--init--标记的代码单元格
init_code_snippets = []
for cell in notebook_json["cells"]:
    if cell["cell_type"] != "code":
        continue
    # 检查单元格是否以#--init--开头（可根据需求调整判断逻辑）
    if cell["source"] and cell["source"][0].strip() == "#--init--":
        init_code_snippets.append("".join(cell["source"]))

# 将收集到的初始化代码写入临时文件
with open("/content/auto_init.py", "w") as f:
    f.write("\n\n".join(init_code_snippets))

# 输出提示并重启运行时
print("即将重启运行时，重启后将自动执行所有初始化单元格")
runtime.restart()

2. 设置重启后自动执行初始化代码 #

创建另一个代码单元格，设置为启动时运行（右键单元格 → 选择"Run on startup"），用来加载并执行临时文件中的初始化代码：

import os

init_file_path = "/content/auto_init.py"
if os.path.exists(init_file_path):
    print("开始执行初始化单元格代码...")
    # 执行初始化代码
    exec(open(init_file_path).read())
    # 执行完成后删除临时文件，避免重复执行
    os.remove(init_file_path)

3. 标记你的初始化单元格 #

给需要作为初始化的代码单元格添加标记：在单元格第一行写入#--init--，比如：

#--init--
!pip install torch torchvision
import numpy as np
import torch
# 其他初始化操作...

注意事项 #

• 确保#--init--是单元格的第一行（如果需要更宽松的匹配，可以修改脚本中的判断逻辑，比如检查单元格内是否包含该标记）
• 初始化代码需要独立可运行，不能依赖其他非初始化单元格的变量（因为重启后运行时会清空所有状态）
• 如果不需要每次重启都执行，可以手动删除临时文件或者取消"Run on startup"的设置

内容的提问来源于stack exchange，提问作者Yaroslav Bulatov