这段内容介绍了将使用一个气球分割数据集(仅包含一个类别:气球)来训练一个气球分割模型。训练过程将以一个预训练在COCO数据集上的模型为基础,这些模型可以在Detectron2的模型库中获取。需要注意的是,COCO数据集本身并不包含“气球”这一类别。通过这一训练过程,我们能够在几分钟内识别这一新类别。
- 自定义数据集数据处理和可视化(通过
dataset_utils.py
)。 - 训练模型(通过
train_balloon.py
)。 - 在验证集上进行推理并可视化结果(通过
validate_balloon.py
)。 - 对模型在验证集上的性能进行评估(通过
evaluation_balloon.py
)。 - 泛化测试推理代码(通过
infer_balloon.py
)
运行环境
使用Dockerfile构建一个包含NVIDIA的PyTorch和Detectron2的镜像
自定义数据集数据处理和可视化
数据集的准备,具体步骤如下:
-
下载数据集
命令行使用wget
下载数据集。wget
是一种网络文件下载工具。wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
-
解压数据集
使用unzip
解压下载的.zip文件,这样我们就可以访问其中的数据。unzip balloon_dataset.zip > /dev/null
通过这段内容,我们可以下载和解压气球数据集,为接下来的模型训练步骤做准备。在后续步骤中,将会解释如何将下载并解压的数据转换为Detectron2的格式,并进行模型训练。
数据集数据处理和可视化
1. 简介
此部分内容展示了如何将气球数据集注册到Detectron2中,并确保其格式符合Detectron2的标准。此外,还演示了如何可视化数据集中的样本以确保注释的正确性。
2. 数据集注册流程
2.1 自定义格式解析
由于数据集是自定义格式,因此需要编写一个函数来解析并将其转换为Detectron2的标准格式。以下是解析并转换数据集的步骤:
气球数据集格式解释
下面是一个样本数据的格式解释:
{
"34020010494_e5cb88e1c4_k.jpg1115004": {
"fileref": "",
"size": 1115004,
"filename": "34020010494_e5cb88e1c4_k.jpg",
"base64_img_data": "",
"file_attributes": {},
"regions": {
"0": {
"shape_attributes": {
"name": "polygon",
"all_points_x": [ ... 点的 x 坐标 ... ],
"all_points_y": [ ... 点的 y 坐标 ... ]
},
"region_attributes": {}
}
}
}
}
字段解释
"fileref"
: 通常为空,表示文件引用。"size"
: 图像文件的大小,以字节为单位。"filename"
: 图像文件的名称。"base64_img_data"
: 图像数据的 Base64 编码,通常为空。"file_attributes"
: 与文件相关的其他属性,通常为空。"regions"
: 包含图像中每个分割区域的注释。- 每个区域有一个唯一的数字键,例如
"0"
. "shape_attributes"
: 包含形状的详细信息。"name"
: 形状的类型,例如此处是"polygon"
。"all_points_x"
: 多边形所有点的 x 坐标数组。"all_points_y"
: 多边形所有点的 y 坐标数组。
"region_attributes"
: 此区域的其他属性,通常为空。
- 每个区域有一个唯一的数字键,例如
转换代码解释
该代码将上述数据集格式转换为 Detectron2 所需的标准格式。
from detectron2.structures import BoxMode
import os
import json
import numpy as np
import cv2
def get_balloon_dicts(img_dir):
json_file = os.path.join(img_dir, "via_region_data.json") # 找到 JSON 文件的路径
with open(json_file) as f:
imgs_anns = json.load(f) # 加载 JSON 文件
dataset_dicts = []
for idx, v in enumerate(imgs_anns.values()): # 遍历所有图像的注释
record = {}
filename = os.path.join(img_dir, v["filename"]) # 获取图像文件路径
height, width = cv2.imread(filename).shape[:2] # 获取图像尺寸
record["file_name"] = filename # 文件名
record["image_id"] = idx # 图像ID
record["height"] = height # 高度
record["width"] = width # 宽度
annos = v["regions"] # 获取注释的 regions
objs = []
for _, anno in annos.items(): # 遍历每个区域注释
assert not anno["region_attributes"] # 确保 region_attributes 为空
anno = anno["shape_attributes"] # 获取形状属性
px = anno["all_points_x"] # 获取所有点的 x 坐标
py = anno["all_points_y"] # 获取所有点的 y 坐标
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)] # 创建 polygon 的顶点坐标
poly = [p for x in poly for p in x] # 将多边形顶点坐标展开为一维列表
obj = {
"bbox": [np.min(px), np.min(py), np.max(px), np.max(py)], # 计算边界框
"bbox_mode": BoxMode.XYXY_ABS, # 边界框的模式,即绝对坐标
"segmentation": [poly], # 分割的多边形
"category_id": 0, # 类别ID,对于气球数据集,这里只有一个类别
}
objs.append(obj) # 将对象添加到对象列表中
record["annotations"] = objs # 将注释添加到记录中
dataset_dicts.append(record) # 将记录添加到数据集字典中
return dataset_dicts # 返回数据集字典
细节解释
-
加载 JSON 文件
json_file = os.path.join(img_dir, "via_region_data.json") with open(json_file) as f: imgs_anns = json.load(f)
这部分代码加载包含注释信息的 JSON 文件。
-
处理每个图像的注释
for idx, v in enumerate(imgs_anns.values()): record = {} filename = os.path.join(img_dir, v["filename"]) height, width = cv2.imread(filename).shape[:2] record["file_name"] = filename record["image_id"] = idx record["height"] = height record["width"] = width
每个图像注释包括文件路径、图像 ID、图像高度和宽度等信息。
-
处理每个区域的注释
annos = v["regions"] objs = [] for _, anno in annos.items(): assert not anno["region_attributes"] anno = anno["shape_attributes"] px = anno["all_points_x"] py = anno["all_points_y"] poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)] poly = [p for x in poly for p in x] obj = { "bbox": [np.min(px), np.min(py), np.max(px), np.max(py)], "bbox_mode": BoxMode.XYXY_ABS, "segmentation": [poly], "category_id": 0, } objs.append(obj)
处理每个区域:读取形状属性(如多边形点),计算边界框,定义分割多边形,并添加注释的详细信息。
-
将注释添加到记录中
record["annotations"] = objs dataset_dicts.append(record)
最后,将处理好的注释添加到记录中,并将记录添加到数据集合中。
通过以上步骤,将自定义的注释格式转换为Detectron2所需的标准数据格式,从而可以进行进一步的模型训练和评估。
2.2 注册数据集
将解析后的数据集注册到Detectron2,这里分别为训练和验证数据集进行注册:
from detectron2.data import DatasetCatalog, MetadataCatalog
for d in ["train", "val"]:
DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("balloon_train")
3. 数据集可视化
为了确保数据集格式正确,可以随机选择一些样本进行可视化:
import random
from detectron2.utils.visualizer import Visualizer
import matplotlib.pyplot as plt
dataset_dicts = get_balloon_dicts("balloon/train")
for d in random.sample(dataset_dicts, 3): # 随机选取三个样本
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(d)
plt.imshow(out.get_image()[:, :, ::-1])
plt.show()
通过以上步骤,我们成功地将自定义格式的气球数据集转换并注册到Detectron2中,并验证其注释的正确性。这样我们就可以通过Detectron2进行模型训练和评估。
4. 完整代码
dataset_utils.py
import os
import json
import numpy as np
import cv2
import random
import argparse
from detectron2.structures import BoxMode
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.utils.visualizer import Visualizer
def get_balloon_dicts(img_dir):
json_file = os.path.join(img_dir, "via_region_data.json")
with open(json_file) as f:
imgs_anns = json.load(f) # Load JSON file
dataset_dicts = []
for idx, v in enumerate(imgs_anns.values()): # Iterate through all images' annotations
record = {}
filename = os.path.join(img_dir, v["filename"]) # Image file path
height, width = cv2.imread(filename).shape[:2] # Get image dimensions
record["file_name"] = filename
record["image_id"] = idx
record["height"] = height
record["width"] = width
annos = v["regions"]
objs = []
for _, anno in annos.items(): # Iterate through each region annotation
assert not anno["region_attributes"]
anno = anno["shape_attributes"]
px = anno["all_points_x"]
py = anno["all_points_y"]
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)] # All points of the shape
poly = [p for x in poly for p in x] # Flatten poly coordinates
obj = {
"bbox": [np.min(px), np.min(py), np.max(px), np.max(py)], # Bounding box
"bbox_mode": BoxMode.XYXY_ABS,
"segmentation": [poly], # Segmentation polygons
"category_id": 0, # Single class "balloon"
}
objs.append(obj)
record["annotations"] = objs
dataset_dicts.append(record)
return dataset_dicts
def register_datasets(dataset_path):
for d in ["train", "val"]:
DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts(os.path.join(dataset_path, d)))
MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
def create_visualizations(dataset_path, output_dir, num_visualizations=3):
os.makedirs(output_dir, exist_ok=True)
dataset_dicts = get_balloon_dicts(os.path.join(dataset_path, "train"))
balloon_metadata = MetadataCatalog.get("balloon_train")
for i, d in enumerate(random.sample(dataset_dicts, num_visualizations)):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(d)
out_file = os.path.join(output_dir, f"visualized_sample_{i}.jpg")
cv2.imwrite(out_file, out.get_image()[:, :, ::-1])
print(f"Saved visualized sample {i} to {out_file}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Visualize and check the balloon dataset")
parser.add_argument("--dataset-path", required=True, help="Path to the root directory of the balloon dataset")
parser.add_argument("--output-dir", required=True, help="Directory to save visualization images")
parser.add_argument("--num-visualizations", type=int, default=3, help="Number of images to visualize for sanity check")
args = parser.parse_args()
register_datasets(args.dataset_path)
create_visualizations(args.dataset_path, args.output_dir, args.num_visualizations)
dataset_utils.py
使用说明
该文件包含:
- 解析和注册气球数据集的功能。
- 可视化数据集样本的功能。
可以通过以下命令运行该文件进行数据集检测和可视化:
python dataset_utils.py --dataset-path /path/to/balloon_dataset --output-dir /path/to/output_directory --num-visualizations 3
模型微调
这段代码演示了如何使用Detectron2框架对COCO预训练的R50-FPN Mask R-CNN模型进行微调,以便在自定义的气球(balloon)数据集上进行训练和评估。
1. 导入依赖库
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2 import model_zoo
import os
这些库提供了Detectron2所需的核心功能:
DefaultTrainer
是一个基于Detectron2的默认训练器,可以极大地简化训练过程。get_cfg
用于获取Detectron2的配置对象。model_zoo
提供了从模型库中获取配置和预训练模型的接口。os
用于文件和目录操作。
2. 获取配置对象并加载预设配置
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
这两行代码执行以下操作:
- 获取一个默认的配置对象
cfg
。 - 从模型库(model zoo)中加载预设的Mask R-CNN配置,并合并到
cfg
中。这些配置文件通常涵盖了模型结构、数据预处理和训练超参数等重要信息。
3. 数据集配置
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()
这里配置了用于训练和测试的数据集:
cfg.DATASETS.TRAIN
指定训练数据集的名称。在这个例子中,它是“balloon_train”。cfg.DATASETS.TEST
被留空,这意味着测试时不使用任何数据集。
4. 数据加载配置
cfg.DATALOADER.NUM_WORKERS = 2
配置数据加载器的工作线程数,NUM_WORKERS=2
意味着会使用两个线程加载数据,通常可以根据硬件条件调节这个参数提高数据加载效率。
5. 模型权重配置
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
设置训练初始权重,这里使用了COCO数据集预训练的Mask R-CNN模型权重。这是迁移学习的一部分,使用预训练权重可以帮助模型更快更好地收敛。
6. 训练超参数配置
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 300
cfg.SOLVER.STEPS = []
这些代码配置了训练的主要超参数:
IMS_PER_BATCH
设置了每批次输入图像的数量,这里设为2,表示每个训练批次使用2张图像。BASE_LR
设置了基础学习率,通常需要进行调试来找到合适的值,这里设为0.00025。MAX_ITER
设置了训练的最大迭代次数。300次迭代通常对这个玩具数据集来说已经足够,但对于实际应用的数据集,通常需要更多的迭代。STEPS
是学习率衰减的步长,这里设为空列表表示不进行学习率衰减。
7. RoI Head配置
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
配置Region of Interest (RoI) head的超参数:
BATCH_SIZE_PER_IMAGE
指定每个图像的RoI batch size。128是一个较快且适用的大多数数据集的设定。NUM_CLASSES
指定数据集中的类别数量,这里只有一个类 “balloon”。
8. 创建输出目录并初始化训练器
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
这行代码确保配置中的输出目录存在,如果不存在,则自动创建。
9. 初始化并启动训练
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
这段代码执行以下操作:
- 使用配置
cfg
实例化一个默认的训练器DefaultTrainer
。 - 使用
resume_or_load
方法。如果resume
参数设为False
,表示从头开始训练。 - 调用
train()
方法启动训练过程。
训练过程
训练过程日志输出详细展示了训练过程中以及训练前的一些设置信息。以下是对输出的详细解释:
模型结构输出
[08/02 09:56:09 d2.engine.defaults]: Model:
GeneralizedRCNN(
(backbone): FPN(
...
(bottom_up): ResNet(
...
)
)
...
)
这部分输出描述了模型的详细结构,从模型的总体概述到各个层级(包括FPN、ResNet和各个卷积操作)的配置。这些信息帮助我们确认模型是否被正确初始化,如以下组件:
- Backbone:使用 FPN(Feature Pyramid Network),并使用预训练的ResNet。
- 头(Heads):包括RPN(Region Proposal Network)、Box Head 和 Mask Head。
数据集分析
[08/02 09:56:10 d2.data.build]: Removed 0 images with no usable annotations. 61 images left.
[08/02 09:56:10 d2.data.build]: Distribution of instances among all 1 categories:
| category | #instances |
|:----------:|:-------------|
| balloon | 255 |
这里输出了数据集的基本信息:
- 总共61张图片(注释文件有效的图片数目)。
- 数据集中每个类别的实例分布,这里只有一个类别 “balloon”,包含255个实例。
训练准备信息
[08/02 09:56:10 d2.data.dataset_mapper]: [DatasetMapper] Augmentations used in training: [ResizeShortestEdge(short_edge_length=(640, 672, 704, 736, 768, 800), max_size=1333, sample_style='choice'), RandomFlip()]
[08/02 09:56:10 d2.data.build]: Using training sampler TrainingSampler
[08/02 09:56:10 d2.data.common]: Serializing the dataset using: <class 'detectron2.data.common._TorchSerializedList'>
[08/02 09:56:10 d2.data.common]: Serializing 61 elements to byte tensors and concatenating them all ...
[08/02 09:56:10 d2.data.common]: Serialized dataset takes 0.17 MiB
[08/02 09:56:10 d2.data.build]: Making batched data loader with batch_size=2
这部分描述了训练前的数据处理与采样设置:
- Data Augmentations:训练使用的增广方法,包括
ResizeShortestEdge
(短边尺寸随机选择)和RandomFlip
(随机翻转)。 - Training Sampler:使用
TrainingSampler
进行数据采样。 - Data Serialization:数据集的序列化方式和大小。
- Batched Data Loader:批数据加载器的创建,
batch_size=2
。
模型权重加载
[08/02 09:56:10 d2.checkpoint.detection_checkpoint]: [DetectionCheckpointer] Loading from https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl ...
model_final_f10217.pkl: 178MB [00:23, 7.71MB/s]
- 权重文件载入:从提供的URL中加载预训练权重文件。
权重形状不匹配警告
Skip loading parameter 'roi_heads.box_predictor.cls_score.weight' to the model due to incompatible shapes: (81, 1024) in the checkpoint but (2, 1024) in the model! You might want to double check if this is expected.
Skip loading parameter 'roi_heads.box_predictor.cls_score.bias' to the model due to incompatible shapes: (81,) in the checkpoint but (2,) in the model! You might want to double check if this is expected.
Skip loading parameter 'roi_heads.box_predictor.bbox_pred.weight' to the model due to incompatible shapes: (320, 1024) in the checkpoint but (4, 1024) in the model! You might want to double check if this is expected.
Skip loading parameter 'roi_heads.box_predictor.bbox_pred.bias' to the model due to incompatible shapes: (320,) in the checkpoint but (4,) in the model! You might want to double check if this is expected.
Skip loading parameter 'roi_heads.mask_head.predictor.weight' to the model due to incompatible shapes: (80, 256, 1, 1) in the checkpoint but (1, 256, 1, 1) in the model! You might want to double check if this is expected.
Skip loading parameter 'roi_heads.mask_head.predictor.bias' to the model due to incompatible shapes: (80,) in the checkpoint but (1,) in the model! You might want to double check if this is expected.
Some model parameters or buffers are not found in the checkpoint:
roi_heads.box_predictor.bbox_pred.{bias, weight}
roi_heads.box_predictor.cls_score.{bias, weight}
roi_heads.mask_head.predictor.{bias, weight}
由于自定义数据集的类别数目与预训练模型的COCO数据集不匹配(COCO有80类,这里只有1类),所以有一些参数形状不匹配,需要跳过这些参数的加载。这是正常现象,可以忽略这些警告。
训练启动
[08/02 09:56:34 d2.engine.train_loop]: Starting training from iteration 0
训练从第0次迭代开始。
训练过程中的日志输出
[08/02 09:56:44 d2.utils.events]: eta: 0:01:57 iter: 19 total_loss: 1.982 loss_cls: 0.622 loss_box_reg: 0.586 loss_mask: 0.6987 loss_rpn_cls: 0.03125 loss_rpn_loc: 0.007351 time: 0.4192 last_time: 0.4716 data_time: 0.0105 last_data_time: 0.0027 lr: 1.6068e-05 max_mem: 2552M
...
[08/02 09:58:46 d2.utils.events]: eta: 0:00:00 iter: 299 total_loss: 0.2945 loss_cls: 0.06335 loss_box_reg: 0.1583 loss_mask: 0.07443 loss_rpn_cls: 0.006397 loss_rpn_loc: 0.004991 time: 0.4299 last_time: 0.3656 data_time: 0.0046 last_data_time: 0.0028 lr: 0.00024917 max_mem: 2765M
在训练过程中,每若干步会输出当前训练状态,包括以下信息:
- eta:预计剩余训练时间。
- iter:当前迭代次数。
- total_loss:总的损失值。
- loss_cls:分类损失。
- loss_box_reg:边界框回归损失。
- loss_mask:掩码损失。
- loss_rpn_cls:RPN分类损失。
- loss_rpn_loc:RPN定位损失。
- time:每步耗时。
- data_time:加载数据的耗时。
- lr:当前学习率。
- max_mem:最大内存使用。
训练结束总结
[08/02 09:58:46 d2.engine.hooks]: Overall training speed: 298 iterations in 0:02:08 (0.4299 s / it)
[08/02 09:58:46 d2.engine.hooks]: Total training time: 0:02:08 (0:00:00 on hooks)
最后,输出了整体训练速度和总的训练时间:
- 每次迭代大约耗时0.4299秒,总耗时2分08秒。
总结
该输出详细记录并展示了训练过程中每个步骤的各类信息,从模型初始化,到数据集处理,再到具体的训练过程。它有助于了解每个环节的执行情况,确认模型结构,检测潜在问题和异常,评估训练效果。
完整代码
### `train_balloon.py`
```python
import os
import argparse
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2 import model_zoo
# Import utility functions from dataset_utils.py
from dataset_utils import register_datasets, create_visualizations
def configure_and_train(dataset_path, output_dir):
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") # Initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2 # Batch size for training
cfg.SOLVER.BASE_LR = 0.00025 # Learning rate
cfg.SOLVER.MAX_ITER = 300 # Number of iterations
cfg.SOLVER.STEPS = [] # No learning rate decay
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # RoIHead batch size
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # Only one class (balloon)
cfg.OUTPUT_DIR = output_dir # Set output directory
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Register and train on balloon dataset with Detectron2")
parser.add_argument("--dataset-path", required=True, help="Path to the root directory of the balloon dataset")
parser.add_argument("--output-dir", required=True, help="Directory to save training outputs")
args = parser.parse_args()
# Register datasets
register_datasets(args.dataset_path)
# Optionally create visualization images (sanity check) before training
create_visualizations(args.dataset_path, args.output_dir)
# Configure and run training
configure_and_train(args.dataset_path, args.output_dir)
该文件是主训练脚本:
- 调用
dataset_utils.py
中注册数据集和进行可视化的函数。 - 配置和启动模型训练。
可以通过以下命令运行训练脚本:
python train_balloon.py --dataset-path /path/to/balloon_dataset --output-dir /path/to/output_directory
查看训练结果
在训练完成后,Detectron2会在–output-dir指定的输出目录中保存各种与训练相关的文件和模型检查点。
目录结构说明
- events.out.tfevents.1722592604.a7863c2c7604.1867.0
- last_checkpoint
- metrics.json
- model_final.pth
详细解释
1. events.out.tfevents.1722592604.a7863c2c7604.1867.0
这是一个TensorFlow事件文件,由 TensorBoard
生成,用于记录训练过程中的标量数据(如损失值、学习率等)。可以使用 TensorBoard
对这些数据进行可视化和分析。
运行 TensorBoard
的命令如下:
tensorboard --logdir result/
然后在浏览器中访问 http://localhost:6006/
,即可查看训练过程中的各种指标。
2. last_checkpoint
这个文件包含了最后一个保存的检查点的路径。通常它只是一行文本,指向最后一个保存的模型文件(在这个例子中,就是 model_final.pth
)。
内容示例:
model_final.pth
这个文件的作用是方便恢复训练时知道从哪一个检查点开始。
3. metrics.json
这个文件记录了训练过程中所有的评估指标(metrics)。内容是一个JSON格式的列表,其中每一对象代表一个步骤或阶段的评估结果,具体字段可能包括迭代步数、损失值、评估指标等。
内容示例:
[
{"iteration": 0, "total_loss": 1.982, "loss_cls": 0.622, "loss_box_reg": 0.586, "loss_mask": 0.6987, "loss_rpn_cls": 0.03125, "loss_rpn_loc": 0.007351, "lr": 1.6068e-05, "time": "2023-08-02T09:56:10"},
{"iteration": 1, "total_loss": 1.830, "loss_cls": 0.5455, "loss_box_reg": 0.6867, "loss_mask": 0.599, "loss_rpn_cls": 0.02135, "loss_rpn_loc": 0.006874, "lr": 3.2718e-05, "time": "2023-08-02T09:57:53"},
...
]
这些记录可以用于后期的分析和可视化,帮助评估模型收敛情况和性能表现。
4. model_final.pth
这个文件是训练过程中保存的最终模型权重。文件格式为 PyTorch .pth
文件,包含了模型的所有参数。
这个文件可以用于模型推断或进一步的微调。
查看训练曲线
TensorBoard 是一个可视化工具,使得研究和调试深度学习训练过程变得直观和简单。
启动 TensorBoard
%tensorboard --logdir output
-
解释:这条命令启动 TensorBoard,并将其指向保存日志的目录(这里是
output
)。 -
参数:
--logdir output
:指定 TensorBoard 日志文件所在的目录。这里output
是日志目录,可以替换为你实际的日志目录路径。
-
作用:启动 TensorBoard 并使其在环境中运行,在新的浏览器标签页或窗口中打开 TensorBoard 界面。
在 TensorBoard 界面中,你可以看到:
- Scalars:包含损失值、学习率、准确度等标量值,你可以查看这些值随时间的变化曲线。
- Graphs:你可以查看模型的计算图,了解模型的结构和数据流。
- Histograms:参数和输出的分布图,可以帮助你理解模型的学习过程。
- Images:训练过程中记录的图像,如输入数据、特征图和生成的掩码等。
- Text 和其他数据类型…
这种可视化可以极大帮助理解模型的训练过程、发现问题并进行调试。TensorBoard 提供了强大的数据可视化和分析功能,是深度学习研究和开发中的重要工具。
使用训练模型在验证集上进行推理和可视化
在 validate_balloon.py
文件中实现使用训练好的模型进行推理并可视化预测结果的功能。以下是完整的代码实现:
validate_balloon.py
import os
import random
import argparse
import cv2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer, ColorMode
from detectron2.data import MetadataCatalog
from dataset_utils import get_balloon_dicts, register_datasets # Import necessary functions from dataset_utils.py
def run_inference_and_visualization(dataset_path, model_weights, output_dir, num_visualizations=3):
# Configuring the model for inference
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only one class (balloon)
cfg.MODEL.WEIGHTS = model_weights # path to the trained model
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold
# Creating predictor
predictor = DefaultPredictor(cfg)
# Registering the validation dataset
register_datasets(dataset_path)
balloon_metadata = MetadataCatalog.get("balloon_val")
# Obtaining validation dataset
dataset_dicts = get_balloon_dicts(os.path.join(dataset_path, "val"))
# Creating output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)
# Run inference and visualize the results
for i, d in enumerate(random.sample(dataset_dicts, num_visualizations)):
img = cv2.imread(d["file_name"])
outputs = predictor(img) # Run inference
v = Visualizer(img[:, :, ::-1],
metadata=balloon_metadata,
scale=0.5,
instance_mode=ColorMode.IMAGE_BW # remove the colors of unsegmented pixels
)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
out_file = os.path.join(output_dir, f"inference_result_{i}.jpg")
cv2.imwrite(out_file, out.get_image()[:, :, ::-1])
print(f"Saved inference result {i} to {out_file}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run inference and visualize the prediction results on balloon dataset")
parser.add_argument("--dataset-path", required=True, help="Path to the root directory of the balloon dataset")
parser.add_argument("--model-weights", required=True, help="Path to the trained model weights")
parser.add_argument("--output-dir", required=True, help="Directory to save visualization images")
parser.add_argument("--num-visualizations", type=int, default=3, help="Number of images to visualize for sanity check")
args = parser.parse_args()
run_inference_and_visualization(args.dataset_path, args.model_weights, args.output_dir, args.num_visualizations)
说明
-
导入和配置:
get_cfg
用于加载配置文件。DefaultPredictor
用于创建推理所需的预测器。Visualizer
和ColorMode
用于可视化推理结果。
-
配置推理模型:
- 使用配置文件
COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml
来设置基本配置。 - 设置
cfg.MODEL.WEIGHTS
为训练好的模型权重路径。 - 设置
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST
为自定义阈值 0.7。
- 使用配置文件
-
注册验证集并获取标注数据:
- 调用
register_datasets(dataset_path)
来注册数据集。 - 获取验证集标注数据。
- 调用
-
进行推理和可视化:
- 遍历验证集中的随机样本图片,使用模型进行推理。
- 使用
Visualizer
绘制实例预测结果并保存到指定输出目录。
-
命令行接口:
- 使用 argparse 模块解析命令行参数,指定数据集路径、模型权重路径、输出目录及可视化图片数量。
使用方法
通过命令行运行该脚本,以在气球验证数据集上运行推理并可视化预测结果:
python validate_balloon.py --dataset-path /path/to/balloon_dataset --model-weights /path/to/model_final.pth --output-dir /path/to/output_directory --num-visualizations 3
以上代码实现了气球数据集的推理和可视化功能。
计算AP(平均精度)指标
创建一个用于评估模型的新脚本 evaluation_balloon.py
。这个脚本将使用 COCO API 中实现的 AP(平均精度)指标来评估训练好的模型。以下是完整的实现:
evaluation_balloon.py
import os
import argparse
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
from detectron2 import model_zoo
from dataset_utils import register_datasets # Import necessary functions from dataset_utils.py
def evaluate_model(dataset_path, model_weights, output_dir):
# Configuring the model for inference
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only one class (balloon)
cfg.MODEL.WEIGHTS = model_weights # path to the trained model
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold
# Creating predictor
predictor = DefaultPredictor(cfg)
# Registering the validation dataset
register_datasets(dataset_path)
# Creating COCO Evaluator
evaluator = COCOEvaluator("balloon_val", output_dir=output_dir)
val_loader = build_detection_test_loader(cfg, "balloon_val")
# Running evaluation
results = inference_on_dataset(predictor.model, val_loader, evaluator)
print(results)
return results
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Evaluate model performance on balloon validation dataset")
parser.add_argument("--dataset-path", required=True, help="Path to the root directory of the balloon dataset")
parser.add_argument("--model-weights", required=True, help="Path to the trained model weights")
parser.add_argument("--output-dir", required=True, help="Directory to save evaluation results")
args = parser.parse_args()
evaluate_model(args.dataset_path, args.model_weights, args.output_dir)
说明
-
导入和配置:
- 从
detectron2.config
中导入get_cfg
用于加载配置文件。 - 从
detectron2.engine
中导入DefaultPredictor
用于创建推理所需的预测器。 - 从
detectron2.evaluation
中导入COCOEvaluator
和inference_on_dataset
用于模型评估。 - 从
detectron2.data
中导入build_detection_test_loader
用于构建测试数据加载器。 - 导入
model_zoo
从 Detectron2 模型库获取预训练模型配置文件。 - 导入
register_datasets
用于注册气球数据集。
- 从
-
配置推理模型:
- 使用 COCO-InstanceSegmentation 配置文件设置基本配置。
- 设置
cfg.MODEL.WEIGHTS
为训练好的模型权重路径。 - 设置
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST
为自定义阈值 0.7。
-
注册验证集并创建评估器:
- 调用
register_datasets(dataset_path)
注册数据集。 - 创建
COCOEvaluator
用于评估模型在验证集上的表现。 - 使用
build_detection_test_loader
构建测试数据加载器。
- 调用
-
进行评估:
- 使用
inference_on_dataset(predictor.model, val_loader, evaluator)
运行评估并打印结果。
- 使用
-
命令行接口:
- 使用 argparse 模块解析命令行参数,指定数据集路径、模型权重路径以及输出目录。
使用方法
通过命令行运行该脚本,以评估模型在气球验证数据集上的表现并输出评估结果:
python evaluation_balloon.py --dataset-path /path/to/balloon_dataset --model-weights /path/to/model_final.pth --output-dir /path/to/output_directory
以上代码实现了模型在气球数据集上的评估功能。
评估结果
看到你的评估结果正常输出了,这说明脚本 evaluation_balloon.py
已经成功运行,并提供了模型在气球验证数据集上的评估结果。
以下是你运行评估后的输出解释:
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.742
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.872
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.806
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.572
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.895
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.232
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.770
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.770
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.635
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.923
这些指标分别显示了不同 IoU 阈值条件下的平均精度(AP)和平均召回率(AR),其中:
- AP 指的是对于所有区域和不同 IoU 阈值下的平均精度,包括 AP50 和 AP75。
- AR 指的是不同情况下的平均召回率。
此外,输出还包括对小(small)、中(medium)、大(large)尺寸目标的评估结果。
最后的 OrderedDict
显示了 bbox
(边界框)和 segm
(分割)结果的详细信息。
OrderedDict([('bbox', {'AP': 74.20085318388486, 'AP50': 87.16228906073731, 'AP75': 80.62367971491025, 'APs': 0.0, 'APm': 57.20849777285421, 'APl': 89.54632560030196}), ('segm', {'AP': 78.19204100062078, 'AP50': 85.07203162782626, 'AP75': 85.07203162782626, 'APs': 0.0, 'APm': 55.7960142168063, 'APl': 96.28293333535034})])
此字典包含了更详细的评估结果:
bbox
:边界框的评估结果。segm
:分割的评估结果。
每个类别的结果中包含不同 IoU 条件下的 AP 值,包括 AP
, AP50
, AP75
, APs
(小尺寸运动对象的AP), APm
(中尺寸运动对象的AP), APl
(大尺寸运动对象的AP)等。
你可以根据这些评估结果,了解模型的性能,并根据需要对模型和数据进行优化。
泛化测试
在网络上随机找一张气球的图片,使用微调后的模型对输入图像进行推理,并将预测的结果进行可视化后保存
infer_balloon.py
import os
import cv2
import argparse
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.utils.visualizer import Visualizer, ColorMode
import warnings
def main(args):
# 定义气球数据集元数据
balloon_metadata = {
"thing_classes": ["balloon"]
}
# 配置推理环境
cfg = get_cfg()
try:
cfg.merge_from_file(args.model_config)
except AssertionError:
# 如果指定的配置文件不存在,则尝试从 model_zoo 加载
print(f"Config file '{args.model_config}' does not exist locally. Attempting to load from model_zoo.")
cfg.merge_from_file(model_zoo.get_config_file(args.model_config))
cfg.MODEL.WEIGHTS = args.model_weights # 使用指定的模型权重文件
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # 设置测试的置信度阈值
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(balloon_metadata["thing_classes"]) # 根据元数据设置类别数
# 创建预测器
predictor = DefaultPredictor(cfg)
# 进行推理和可视化
im = cv2.imread(args.input_image)
outputs = predictor(im) # 获取推理结果
v = Visualizer(im[:, :, ::-1],
metadata=balloon_metadata, # 直接传递元数据
scale=0.5,
instance_mode=ColorMode.IMAGE_BW # 去除未分割像素的颜色
)
out = v.draw_instance_predictions(outputs['instances'].to('cpu'))
# 保存可视化结果
cv2.imwrite(args.output_image, out.get_image()[:, :, ::-1])
print(f"Saved visualized result to {args.output_image}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Inference using a fine-tuned model")
parser.add_argument("--model-config", required=True, help="Path to the model config file or model_zoo name (e.g., COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml)")
parser.add_argument("--model-weights", required=True, help="Path to the model weights file (e.g., model_final.pth)")
parser.add_argument("--input-image", required=True, help="Path to the input image file")
parser.add_argument("--output-image", required=True, help="Path to save the output image with visualized predictions")
args = parser.parse_args()
# 忽略 torch.meshgrid 警告
warnings.filterwarnings("ignore", category=UserWarning, message=r"torch.meshgrid")
main(args)
说明
-
定义气球数据集元数据:
balloon_metadata
是一个字典,其中包含thing_classes
,将类别名称指定为"balloon"
。
-
配置推理环境:
- 从指定的配置文件加载模型配置。
- 加载微调后的模型权重文件。
- 设置模型的类别数量,这里根据
balloon_metadata
中的类别数量进行配置。
-
进行推理和可视化:
- 使用
DefaultPredictor
对输入图像进行推理。 - 使用
Visualizer
对预测结果进行可视化,并将balloon_metadata
作为元数据直接传递给Visualizer
。
- 使用
-
保存可视化结果:
- 使用 OpenCV 将结果图像保存到指定路径。
使用示例
通过以下命令运行脚本,确保路径和参数正确:
python3 infer_balloon.py --model-config ../detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml --model-weights ./result/model_final.pth --input-image ./balloon/val/2917282960_06beee649a_b.jpg --output-image ./ft1.jpg
总结
现在已经具备了完成以下工作的脚本:
- 数据处理和可视化(通过
dataset_utils.py
)。 - 训练模型(通过
train_balloon.py
)。 - 在验证集上进行推理并可视化结果(通过
validate_balloon.py
)。 - 对模型在验证集上的性能进行评估(通过
evaluation_balloon.py
)。 - 泛化测试推理代码(通过
infer_balloon.py
)
这些脚本共同构成了一个完整的机器学习模型训练和验证流程。
运行环境
使用Dockerfile构建一个包含NVIDIA的PyTorch和Detectron2的镜像