Huấn Luyện Trước Mô Hình Phân Vùng Ngữ Nghĩa Trên Tập Dữ Liệu COCO
Là các kỹ sư thị giác máy tính và học sâu, chúng ta thường fine-tune (tinh chỉnh) các mô hình phân vùng ngữ nghĩa cho nhiều tác vụ khác nhau. Để làm điều này, PyTorch cung cấp một số mô hình đã được huấn luyện trước trên tập dữ liệu COCO. Mô hình nhỏ nhất có sẵn trên nền tảng Torchvision là mô hình LRASPP MobileNetV3 với 3.2 triệu tham số. Nhưng điều gì sẽ xảy ra nếu chúng ta muốn mô hình nhỏ hơn nữa? Chúng ta có thể làm được, nhưng chúng ta cũng cần phải huấn luyện trước nó. Bài viết này sẽ giải quyết vấn đề này. Chúng ta sẽ sửa đổi kiến trúc LRASPP để tạo ra một mô hình phân vùng ngữ nghĩa với backbone MobileNetV3 Small. Không chỉ vậy, chúng ta cũng sẽ huấn luyện trước mô hình phân vùng ngữ nghĩa trên tập dữ liệu COCO.
Nhảy đến phần Tải Code
Hình 1. Kết quả sau khi huấn luyện một mô hình phân vùng ngữ nghĩa 1.07 triệu tham số trên tập dữ liệu COCO.
Đến cuối bài viết này, bạn sẽ có một ý tưởng toàn diện về việc tạo ra các mô hình phân vùng ngữ nghĩa nhỏ hơn nhưng vẫn hiệu quả và phương pháp huấn luyện trước chúng.
Những gì chúng ta sẽ đề cập đến trong khi huấn luyện trước mô hình phân vùng ngữ nghĩa trên tập dữ liệu COCO?
- Làm thế nào để sửa đổi code mô hình phân vùng LRASPP chính thức của Torchvision để chuyển nó với backbone theo lựa chọn của chúng ta?
- Làm thế nào để sửa đổi các script huấn luyện trước phân vùng ngữ nghĩa chính thức từ Torchvision?
- Quy trình chạy suy luận trên hình ảnh và video sau khi huấn luyện trước mô hình phân vùng trên tập dữ liệu COCO là gì?
Tại Sao Chúng Ta Cần Huấn Luyện Trước Một Mô Hình Phân Vùng Ngữ Nghĩa Trên Tập Dữ Liệu COCO?
Mục tiêu của bài viết này là hai thứ:
- Tạo ra một mô hình phân vùng ngữ nghĩa nhỏ (khoảng 1.1 triệu tham số) cho các thiết bị có tài nguyên hạn chế.
- Tuy nhiên, các mô hình nhỏ như vậy cũng cần huấn luyện trước trước khi chúng ta có thể fine-tune chúng cho các tác vụ hạ nguồn. Vì lý do này, chúng ta cần phải huấn luyện trước mô hình trên một tập dữ liệu phân vùng ngữ nghĩa lớn trước giai đoạn fine-tuning.
Chúng ta có cần chuyển đổi tập dữ liệu COCO thành một tập dữ liệu phân vùng ngữ nghĩa từ tập dữ liệu phát hiện hoặc phân vùng instance ban đầu của nó không? Không, vì chúng ta sẽ sử dụng các script Torchvision chính thức để huấn luyện trước, trình tải dữ liệu sẽ xử lý việc chuyển đổi các mask từ các file chú thích JSON. Chúng ta chỉ cần tải xuống tập dữ liệu COCO gốc và chỉ định script huấn luyện đến thư mục chính xác.
Chúng Ta Sẽ Tuân Theo Quy Trình Nào Để Huấn Luyện Trước Mô Hình Phân Vùng Ngữ Nghĩa Trên Tập Dữ Liệu COCO?
- Đầu tiên, chúng ta sẽ tải xuống tập dữ liệu COCO chính thức.
- Tiếp theo, chúng ta sẽ thiết lập thư mục code và kiểm tra cấu trúc thư mục.
- Sau đó, chúng ta sẽ xem xét file mô hình phân vùng ngữ nghĩa mới về cách tạo một mô hình LRASPP tùy chỉnh nhỏ và hiệu quả.
- Cuối cùng, chúng ta sẽ bắt đầu quá trình huấn luyện trước.
Tải Xuống Tập Dữ Liệu COCO
Bạn có thể tải xuống tập dữ liệu COCO từ [đây]([Đường dẫn Kaggle COCO Dataset]) trên Kaggle. Sau khi tải xuống và giải nén nội dung của nó, bạn sẽ thấy cấu trúc thư mục sau.
coco2017/
├── annotations
│ ├── captions_train2017.json
│ ├── captions_val2017.json
│ ├── instances_train2017.json
│ ├── instances_val2017.json
│ ├── person_keypoints_train2017.json
│ └── person_keypoints_val2017.json
├── test2017 [40670 entries exceeds filelimit, not opening dir]
├── train2017 [118287 entries exceeds filelimit, not opening dir]
└── val2017 [5000 entries exceeds filelimit, not opening dir]
content_copy download Use code with caution.
Mặc dù tập dữ liệu COCO chủ yếu là một tập dữ liệu phát hiện và phân vùng instance, nó cũng có thể được sử dụng để huấn luyện trước phân vùng ngữ nghĩa. Trên thực tế, tất cả các mô hình phân vùng ngữ nghĩa Torchvision đều được huấn luyện trước trên tập dữ liệu COCO. Tuy nhiên, một sửa đổi nhỏ là thay vì huấn luyện trên 80 lớp COCO, chỉ những hình ảnh tương ứng với 20 lớp phân vùng ngữ nghĩa Pascal VOC mới được xem xét.
Đây là một vài hình ảnh cùng với các mask phân vùng của chúng từ tập dữ liệu COCO.
Hình 2. Các mẫu ground truth phân vùng ngữ nghĩa từ tập dữ liệu COCO.
Cấu Trúc Thư Mục Dự Án
Hãy xem toàn bộ cấu trúc thư mục dự án.
├── coco2017
├── inference_data
│ ├── images
│ │ ├── image_1.jpg
│ │ ├── image_2.jpg
│ │ └── image_3.jpg
│ └── videos
│ ├── video_1.mp4
│ └── video_2.mp4
├── models
│ ├── custom_lraspp_mobilenetv3_small.py
│ └── __init__.py
├── outputs
│ └── lraspp_mobilenetv3small [31 entries exceeds filelimit, not opening dir]
├── coco_utils.py
├── config.py
├── custom_utils.py
├── inference_image.py
├── inference_video.py
├── __init__.py
├── presets.py
├── README.md
├── requirements.txt
├── train.py
├── transforms.py
├── utils.py
└── v2_extras.py
content_copy download Use code with caution.
- Thư mục coco2017 chứa tập dữ liệu mà chúng ta đã tải xuống ở trên.
- Thư mục inference_data chứa các hình ảnh và video mà chúng ta sẽ sử dụng để suy luận sau khi giai đoạn huấn luyện trước hoàn tất.
- Vì chúng ta sẽ tạo một mô hình phân vùng ngữ nghĩa tùy chỉnh, chúng ta sẽ giữ tất cả các file code liên quan trong thư mục models.
- Thư mục outputs chứa các mô hình từ giai đoạn huấn luyện trước.
- Ngay bên trong thư mục dự án mẹ, chúng ta có các script Torchvision chính thức và một vài script tùy chỉnh. Các file config.py, custom_utils.py, inference_image.py và inference_video.py là các script tùy chỉnh và phần còn lại là từ các tham chiếu phân vùng ngữ nghĩa Torchvision.
Tất cả các file code và trọng số đã được huấn luyện trước đều có sẵn thông qua phần tải xuống.
Tải Code
[Tải xuống Source Code cho Hướng Dẫn này]([Đường dẫn tải xuống])
Cài Đặt Dependencies
Bạn có thể cài đặt các yêu cầu bằng file requirements.txt.
pip install -r requirements.txt
content_copy download Use code with caution.
Nếu bạn muốn huấn luyện trước, bạn nên tải xuống tập dữ liệu COCO 2017 và giữ nó trong cấu trúc trên trước khi chuyển tiếp.
Huấn Luyện Trước Một Mô Hình Phân Vùng Ngữ Nghĩa Trên Tập Dữ Liệu COCO
Chúng ta sẽ chủ yếu khám phá việc chuẩn bị mô hình trong phần code. Vì codebase lớn, rất khó để xem xét tất cả các file Python. Hãy thoải mái đi sâu vào bất kỳ file code nào để có được sự hiểu biết sâu hơn về dự án.
Mô Hình Phân Vùng Ngữ Nghĩa LRASPP MobileNetV3
Chúng ta sẽ chuyển đổi code nguồn chính thức từ Torchvision để xây dựng mô hình phân vùng ngữ nghĩa LRASPP. Chúng ta sẽ chuyển đổi MobileNetV3 Large backbone với MobileNetV3 Small backbone.
Code có trong file models/custom_lraspp_mobilenetv3_small.py. Khối sau chứa tất cả code mà chúng ta cần.
"""
Modified from URL: https://pytorch.org/vision/stable/_modules/torchvision/models/segmentation/lraspp.html#LRASPP_MobileNet_V3_Large_Weights
"""
from torchvision.models.mobilenet import MobileNetV3
from torchvision.models.segmentation import LRASPP
from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights
from torchvision.models._utils import IntermediateLayerGetter, _ovewrite_value_param
from typing import Optional, Any
from torchinfo import summary
def _lraspp_mobilenetv3(backbone: MobileNetV3, num_classes: int) -> LRASPP:
backbone = backbone.features
# Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
# The first and last blocks are always included because they are the C0 (conv1) and Cn.
stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1]
low_pos = stage_indices[-4] # use C2 here which has output_stride = 8
high_pos = stage_indices[-1] # use C5 which has output_stride = 16
low_channels = backbone[low_pos].out_channels
high_channels = backbone[high_pos].out_channels
backbone = IntermediateLayerGetter(backbone, return_layers={str(low_pos): "low", str(high_pos): "high"})
return LRASPP(backbone, low_channels, high_channels, num_classes)
def lraspp_mobilenet_v3_small(
*,
weights = None,
progress: bool = True,
num_classes: Optional[int] = None,
weights_backbone: Optional[MobileNet_V3_Small_Weights] = MobileNet_V3_Small_Weights.IMAGENET1K_V1,
**kwargs: Any,
) -> LRASPP:
"""Constructs a Lite R-ASPP Network model with a MobileNetV3-Large backbone from
`Searching for MobileNetV3 <https://arxiv.org/abs/1905.02244>`_ paper.
Args:
weights (:class:`~torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights`, optional): The
pretrained weights to use. See
:class:`~torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights` below for
more details, and possible values. By default, no pre-trained
weights are used.
progress (bool, optional): If True, displays a progress bar of the
download to stderr. Default is True.
num_classes (int, optional): number of output classes of the model (including the background).
aux_loss (bool, optional): If True, it uses an auxiliary loss.
weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained
weights for the backbone.
**kwargs: parameters passed to the ``torchvision.models.segmentation.LRASPP``
base class. Please refer to the `source code
<https://github.com/pytorch/vision/blob/main/torchvision/models/segmentation/lraspp.py>`_
for more details about this class.
"""
if kwargs.pop("aux_loss", False):
raise NotImplementedError("This model does not use auxiliary loss")
weights_backbone = MobileNet_V3_Small_Weights.verify(weights_backbone)
if weights is not None:
print('Loading LRASPP segmentation pretrained weights...')
weights_backbone = None
num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
elif num_classes is None:
print('Loading custom backbone')
num_classes = 21
backbone = mobilenet_v3_small(weights=weights_backbone, dilated=True)
model = _lraspp_mobilenetv3(backbone, num_classes)
if weights is not None:
print('Loading LRASPP segmentation pretrained weights...')
model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
return model
if __name__ == '__main__':
model = lraspp_mobilenet_v3_small()
summary(model)
Trong khối code trên, hàm _lraspp_mobilenetv3 trích xuất các đặc trưng từ low_channels và high_channels. Chúng được sử dụng để xây dựng LRASPP decoder head bằng cách hợp nhất các đặc trưng. Để có kết quả tốt nhất, việc lựa chọn các lớp một cách chiến lược có thể quan trọng vì các backbone khác nhau học các loại đặc trưng khác nhau ở các độ sâu khác nhau.
Hàm lraspp_mobilenet_v3_small xây dựng toàn bộ mô hình. Nó chấp nhận ánh xạ trọng số backbone là MobileNet_V3_Small_Weights.IMAGENET1K_V1 trong trường hợp của chúng ta. Vì chúng ta không tải trọng số đã được huấn luyện trước của một mô hình phân vùng ngữ nghĩa hoàn chỉnh, mô hình cuối cùng được tạo bằng cách sử dụng backbone đã được huấn luyện trước và gọi hàm _lraspp_mobilenetv3.
Thực thi file này cho chúng ta kết quả sau.
python models/custom_lraspp_mobilenetv3_small.py
Loading custom backbone
===========================================================================
Layer (type:depth-idx) Param #
===========================================================================
LRASPP --
├─IntermediateLayerGetter: 1-1 --
│ └─Conv2dNormActivation: 2-1 --
│ │ └─Conv2d: 3-1 432
│ │ └─BatchNorm2d: 3-2 32
│ │ └─Hardswish: 3-3 --
│ └─InvertedResidual: 2-2 --
│ │ └─Sequential: 3-4 744
│ └─InvertedResidual: 2-3 --
│ │ └─Sequential: 3-5 3,864
│ └─InvertedResidual: 2-4 --
│ │ └─Sequential: 3-6 5,416
│ └─InvertedResidual: 2-5 --
│ │ └─Sequential: 3-7 13,736
│ └─InvertedResidual: 2-6 --
│ │ └─Sequential: 3-8 57,264
│ └─InvertedResidual: 2-7 --
│ │ └─Sequential: 3-9 57,264
│ └─InvertedResidual: 2-8 --
│ │ └─Sequential: 3-10 21,968
│ └─InvertedResidual: 2-9 --
│ │ └─Sequential: 3-11 29,800
│ └─InvertedResidual: 2-10 --
│ │ └─Sequential: 3-12 91,848
│ └─InvertedResidual: 2-11 --
│ │ └─Sequential: 3-13 294,096
│ └─InvertedResidual: 2-12 --
│ │ └─Sequential: 3-14 294,096
│ └─Conv2dNormActivation: 2-13 --
│ │ └─Conv2d: 3-15 55,296
│ │ └─BatchNorm2d: 3-16 1,152
│ │ └─Hardswish: 3-17 --
├─LRASPPHead: 1-2 --
│ └─Sequential: 2-14 --
│ │ └─Conv2d: 3-18 73,728
│ │ └─BatchNorm2d: 3-19 256
│ │ └─ReLU: 3-20 --
│ └─Sequential: 2-15 --
│ │ └─AdaptiveAvgPool2d: 3-21 --
│ │ └─Conv2d: 3-22 73,728
│ │ └─Sigmoid: 3-23 --
│ └─Conv2d: 2-16 525
│ └─Conv2d: 2-17 2,709
===========================================================================
Total params: 1,077,954
Trainable params: 1,077,954
Non-trainable params: 0
===========================================================================
Mô hình cuối cùng với 21 lớp (đại diện cho các lớp tập dữ liệu Pascal VOC) chứa 1.07 triệu tham số.
Nếu bạn muốn tìm hiểu thêm về việc sửa đổi mô hình phân vùng Torchvision, bạn có thể xem bài viết này, nơi chúng ta xây dựng mô hình DeepLabV3 tùy chỉnh.
Huấn Luyện Mô Hình LRASPP MobileNetV3 Small Tùy Chỉnh
Chúng ta sẽ bỏ qua phần còn lại của code để ngắn gọn trong bài viết này. Sửa đổi duy nhất khác mà chúng ta đã thực hiện trong code là trong script train.py. Tại đây, chúng ta đã comment việc tải mô hình Torchvision và tải mô hình của riêng mình.
# model = torchvision.models.get_model(
# args.model,
# weights=args.weights,
# weights_backbone=args.weights_backbone,
# num_classes=num_classes,
# aux_loss=args.aux_loss,
# )
model = lraspp_mobilenet_v3_small(num_classes=num_classes)
summary(model)
Kết quả huấn luyện được hiển thị ở đây được thực hiện trên một máy có GPU RTX 3080 10GB, CPU i7 thế hệ thứ 10 và 32 GB RAM.
Chúng ta có thể thực thi lệnh sau để bắt đầu huấn luyện.
torchrun --nproc_per_node=1 train.py --lr 0.02 --dataset coco -b 64 -j 8 --amp --output-dir outputs/lraspp_mobilenetv3small --data-path coco2017
Đây là các đối số dòng lệnh mà chúng ta sử dụng:
- –nproc_per_node=1: Vì chúng ta đang huấn luyện trên một máy duy nhất, số lượng node cho chúng ta là 1.
- –lr: Đây là learning rate cơ bản cho optimizer.
- –dataset: Tập dữ liệu huấn luyện trước. Vì chúng ta đang huấn luyện trên tập dữ liệu COCO, giá trị ở đây là coco.
- -b: Batch size cho các trình tải dữ liệu.
- -j: Điều này chỉ định số lượng worker song song cho các trình tải dữ liệu.
- –amp: Đây là một đối số boolean chỉ định script huấn luyện sử dụng huấn luyện độ chính xác hỗn hợp.
- –output_dir: Đường dẫn nơi các trọng số đã huấn luyện sẽ được lưu trữ.
- –data-path: Đường dẫn thư mục chứa tập dữ liệu COCO.
Theo mặc định, quá trình huấn luyện chạy trong 30 epoch, chúng ta có thể thay đổi bằng đối số –epochs.
Khối sau hiển thị các đầu ra từ epoch cuối cùng.
Epoch: [29] [ 0/1445] eta: 1:14:42 lr: 0.000936160425031581 loss: 0.3725 (0.3725) time: 3.1019 data: 2.8259 max mem: 8356
Epoch: [29] [ 10/1445] eta: 0:14:57 lr: 0.0009303236046314702 loss: 0.3079 (0.2989) time: 0.6255 data: 0.3819 max mem: 8356
Epoch: [29] [ 20/1445] eta: 0:12:02 lr: 0.0009244827124700352 loss: 0.2815 (0.2971) time: 0.3771 data: 0.1361 max mem: 8356
Epoch: [29] [ 30/1445] eta: 0:10:41 lr: 0.0009186377170825067 loss: 0.2815 (0.2914) time: 0.3585 data: 0.1244 max mem: 8356
.
.
.
Epoch: [29] [1430/1445] eta: 0:00:05 lr: 1.4429624715447326e-05 loss: 0.2976 (0.2985) time: 0.3901 data: 0.1544 max mem: 8356
Epoch: [29] [1440/1445] eta: 0:00:01 lr: 4.672978643683746e-06 loss: 0.3055 (0.2986) time: 0.3240 data: 0.1058 max mem: 8356
Epoch: [29] Total time: 0:08:54
Test: [ 0/5000] eta: 2:49:24 time: 2.0329 data: 1.4863 max mem: 8356
Test: [ 100/5000] eta: 0:02:25 time: 0.0100 data: 0.0020 max mem: 8356
.
.
.
Test: [4800/5000] eta: 0:00:01 time: 0.0086 data: 0.0017 max mem: 8356
Test: [4900/5000] eta: 0:00:00 time: 0.0083 data: 0.0017 max mem: 8356
Test: Total time: 0:00:43
global correct: 89.8
average row correct: ['93.4', '74.5', '67.0', '64.6', '51.0', '21.4', '73.0', '51.2', '86.9', '22.4', '71.8', '69.4', '77.1', '71.9', '75.6', '84.6', '38.1', '74.2', '48.3', '80.4', '53.2']
IoU: ['88.7', '60.3', '54.6', '47.4', '42.7', '18.4', '66.7', '40.9', '69.7', '19.3', '59.9', '33.4', '61.8', '59.3', '62.1', '73.3', '23.5', '57.7', '40.5', '68.9', '43.5']
mean IoU: 52.0
Training time 5:00:17
Mô hình đạt mean IoU là 52% ở epoch cuối cùng. Chúng ta sẽ sử dụng mô hình này để chạy suy luận.
Chạy Suy Luận Sử Dụng Mô Hình Phân Vùng Tùy Chỉnh Đã Được Huấn Luyện Trước Trên Tập Dữ Liệu COCO
Có hai script cho suy luận:
- inference_image.py: Để chạy suy luận trên hình ảnh.
- inference_video.py: Để chạy suy luận trên video.
Hãy bắt đầu với suy luận hình ảnh.
python inference_image.py --input inference_data/images/ --model outputs/lraspp_mobilenetv3small/model_29.pth
content_copy download Use code with caution.
Ở đây, chúng ta chuyển các đường dẫn đến thư mục hình ảnh đầu vào và các trọng số mô hình. Hầu hết các hình ảnh đều chứa con người để kiểm tra xem mô hình phân vùng ngữ nghĩa tùy chỉnh nhỏ của chúng ta có thể phân vùng những người một cách thích hợp hay không.
Sau đây là các kết quả.
Hình 3. Kết quả suy luận hình ảnh sau khi huấn luyện trước mô hình phân vùng ngữ nghĩa trên tập dữ liệu COCO.
Kết quả không hoàn hảo nhưng không tệ đối với một mô hình 1.07M tham số. Trong hầu hết các trường hợp, nó có thể phân vùng đầu ra của những người một cách chính xác.
Chạy suy luận trên video sẽ cho chúng ta một ý tưởng tốt hơn về khả năng của mô hình.
python inference_video.py --input inference_data/videos/video_1.mp4 --model outputs/lraspp_mobilenetv3small/model_29.pth
Đối với suy luận video, chúng ta cung cấp đường dẫn đến một file video.
Video 1. Phân vùng người bằng mô hình phân vùng ngữ nghĩa đã được huấn luyện trước.
Trong trường hợp này, những người đang di chuyển và mô hình dường như có thể phân vùng chúng khá tốt. Hơn nữa, chúng ta đang nhận được trung bình 105 FPS trên GPU RTX 3080.
Mặc dù đáng chú ý rằng trong khi chạy suy luận, chúng ta đang thay đổi kích thước các frame thành 416×512 (chiều cao x chiều rộng). Một số tăng tốc đến từ đó.
Bây giờ, hãy chạy mô hình trên một cảnh phức tạp.
python inference_video.py --input inference_data/videos/video_3.mp4 --model outputs/lraspp_mobilenetv3small/model_29.pth
Video 2. Sử dụng mô hình phân vùng đã được huấn luyện trước để phân vùng các đối tượng trong một cảnh giao thông phức tạp.
Kết quả không hoàn hảo ở đây, tuy nhiên, chúng cũng không quá tệ. Với nhiều huấn luyện và tăng cường dữ liệu hơn, mô hình có thể được cải thiện rất nhiều. Nó đang gặp khó khăn trong việc phân vùng những người ở xa.
Tóm Tắt và Kết Luận
Trong bài viết này, chúng ta đã sửa đổi mô hình Torchvision LRASPP với một backbone MobileNetV3 tùy chỉnh và thực hiện huấn luyện trước trên tập dữ liệu COCO. Điều này cung cấp cho chúng ta kiến thức và ý tưởng để sửa đổi các backbone trong các mô hình phân vùng Torchvision được xác định trước. Mô hình này, nhỏ và cung cấp FPS lớn hơn thời gian thực, có thể được sử dụng để fine-tuning cụ thể cho các tác vụ tiếp theo. Tôi hy vọng rằng bài viết này đáng giá thời gian của bạn.
Nếu bạn có bất kỳ nghi ngờ, suy nghĩ hoặc đề xuất nào, vui lòng để lại chúng trong phần bình luận. Tôi chắc chắn sẽ giải quyết chúng.
Bạn có thể liên hệ với tôi bằng phần Liên Hệ. Bạn cũng có thể tìm thấy tôi trên LinkedIn và Twitter.
