如何将本地训练的TensorFlow及Magenta MusicVAE模型导入Google Colab

步骤1：整理本地训练完成的模型文件 #

你在本地训练生成的MusicVAE模型，需先将完整的训练输出目录打包，目录核心包含以下文件：

• 三类checkpoint核心文件：model.ckpt-xxx.data-00000-of-00001、model.ckpt-xxx.index、model.ckpt-xxx.meta
• 训练时使用的.yaml格式配置文件，记录模型结构、参数设置
• 可选：若你已导出冻结后的*.pb格式推理模型，可一并打包
  打包为zip或tar.gz格式即可。

步骤2：将模型压缩包上传至Colab可访问路径 #

有两种常用上传方案：

• 方案1：直接上传至Colab临时存储
  在Colab中运行以下代码，触发文件上传控件，选择本地的模型压缩包即可，上传后文件保存在Colab当前工作目录：

  from google.colab import files
  uploaded = files.upload()
  

• 方案2：通过Google Drive挂载（推荐，避免Colab会话重置后文件丢失）
  先将压缩包上传到你的Google Drive任意目录，再在Colab中运行挂载代码，按提示完成授权即可访问Drive内所有文件：

  from google.colab import drive
  drive.mount('/content/drive')
  

  挂载完成后，Google Drive根目录对应Colab路径为/content/drive/MyDrive。

步骤3：环境配置与模型加载 #

3.1 安装Magenta依赖 #

Colab默认未预装Magenta框架，先运行安装命令：

!pip install magenta

注意：请保证Colab安装的Magenta版本与你本地训练时的版本一致，避免参数不兼容报错，可指定版本安装，例如!pip install magenta==2.1.4

3.2 解压模型包 #

将上传的压缩包解压到指定路径，例如：

!unzip /content/drive/MyDrive/my_musicvae_model.zip -d /content/musicvae_model

3.3 加载模型 #

运行以下代码加载自定义模型，注意替换为你自己的配置名和checkpoint路径：

import magenta.music as mm
from magenta.models.music_vae import configs
from magenta.models.music_vae.trained_model import TrainedModel

# 替换为你训练时使用的配置名称，或读取本地yaml配置文件
config = configs.CONFIG_MAP['你的自定义配置名']
# 替换为解压后的checkpoint前缀，例如/model/path/model.ckpt-10000
# 若你使用冻结后的pb模型，直接传入pb文件的完整路径即可
model = TrainedModel(config, batch_size=4, checkpoint_dir_or_path='/content/musicvae_model/model.ckpt-xxx')

3.4 验证加载结果 #

可运行以下代码生成测试旋律，确认模型加载正常：

# 生成4个8小节旋律片段
samples, _, _ = model.sample(n=4, length=32)
# 导出为MIDI文件验证
for idx, sample in enumerate(samples):
    mm.sequence_proto_to_midi_file(sample, f'test_output_{idx}.mid')

内容的提问来源于stack exchange，提问作者Sravya Alla