在Google Colab中保存和加载模型权重

Google Colab是一个强大的在线平台，它提供了预装了许多流行Python包的环境，非常适合进行模型实验。因此，Roboflow模型库包含了许多可在Google Colab上使用的免费开源计算机视觉模型。然而，Colab也有一定的限制，比如分配的计算资源仅限于12小时。如果浏览器会话检测到一小时的不活动（例如用户不在该标签页中活跃），内核将断开连接。因此，保存模型的权重并在未来的会话中重新加载它们非常重要。保存模型权重意味着保存模型训练后的拟合度。重新加载模型权重意味着在未来的实验中使用这些保存的权重——即使这是Colab中的一个新会话。

在本文中，将介绍如何在Roboflow模型库中的笔记本中保存和重新加载YOLOv5模型的权重，但这些步骤适用于任何Colab会话（需要修改文件路径）。首先，需要有想要保存的权重以及想要保存这些权重的目标位置。在训练模型后，模型的权重被存储为Colab会话中的一个文件。在YOLOv5笔记本示例中，这些权重保存在runs文件夹中。具体来说，文件路径是：/content/yolov5/runs/exp0_yolov5s_results/weights/best.pt，其中权重被称为best.pt。（专业提示：要确定Colab中任何文件的文件路径，可以使用左侧的导航窗格，右键单击给定文件，并选择“复制文件路径”。）

一旦有了权重文件的文件路径，可以将此文件保存在本地或Google Drive中。建议将权重保存到Google Drive中。如何在Google Colab中下载文件？这很简单：使用以下代码片段，将example.txt替换为权重文件的路径。还可以将文件下载并保存到Google Drive中，这可能更方便，因为可能希望将权重文件与Google帐户一起使用。如何将文件从Google Colab下载到Google Drive？有两个步骤。首先，通过运行以下代码将Google Drive连接到Google Colab会话。这将提示访问一个单独的页面并复制/粘贴授权码。

如何将Google Drive链接到Google Colab笔记本。其次，将文件从Google Colab笔记本复制到Google Drive。这需要指定（1）想要复制的文件的路径（在这种情况下是权重）和（2）正在Google Drive中保存权重的位置。在下面的代码片段中，假设使用的是YOLOv5笔记本模型权重位置，并且正在将它们保存到Google Drive的根目录。

将权重从YOLOv5笔记本复制到Google Drive。运行此单元格后，权重将作为名为best.pt的文件在Google Drive中可用。一旦保存了权重，可能希望在稍后的Google Colab会话中重新加载它们。加载模型权重需要已经构建了模型架构和环境。但是，不会使用该架构来训练新模型，而是使用该架构对新数据集进行推理。在例子中，将再次基于YOLOv5笔记本进行图像识别和对象检测。（在YOLOv5笔记本中，在调用detect.py脚本时进行推理。）

要将权重添加到Colab会话中，可以从本地计算机或Google Drive加载权重。需要知道模型权重的现有位置，以及在Colab笔记本中添加它们的位置。将介绍如何从本地计算机和Google Drive加载文件。如何在Google Colab中上传文件？对于本地计算机上的文件，将使用Google Colab中的用户界面。在左上角，选择文件夹图标以查看给定目录中可用的文件夹列表。在这里，有一个上传文件的图标。选择此图标，然后选择本地计算机上权重的位置。（请注意，此上传图标将文件从本地目录加载到选择上传图标的位置。）

如何在Google Colab中上传文件？首先，使用以下代码片段将Google Drive和Google Colab笔记本链接起来（这将需要一个授权码）。现在，可以将权重从Google Drive加载到Colab笔记本中。这需要知道权重在Google Drive中的文件路径，以及想要将权重复制到Colab会话中的哪个位置。假设权重位于Google Drive的基础目录中（这是在本博客文章中的先前示例中保存它们的位置）。还将假设将模型权重加载到称为weights的YOLO5目录结构中的目录中。要将权重（称为best.pt）从Google Drive复制到Colab笔记本的weights目录，将在单元格中运行此操作：

从Google Drive将权重复制到YOLO5目录中的名为“weights”的文件夹。现在已经在Colab会话中有了模型权重，可以使用它们与模型一起使用。这需要运行所有设置模型的单元格（在YOLOv5笔记本中，这是在运行train.py单元格之前的所有内容）。需要知道权重文件的文件路径，并将该文件路径传递给运行推理的模型。在上面的YOLOv5示例中，将权重称为best.pt加载到content/yolov5/weights目录中。因此，使用模型权重调用推理的脚本应该引用这个权重文件路径。它看起来像这样：