DINOv2 cho Phân Vùng Ngữ Nghĩa
Việc huấn luyện các mô hình phân vùng ngữ nghĩa (semantic segmentation) thường tốn thời gian và đòi hỏi nhiều tài nguyên tính toán. Tuy nhiên, với các backbone DINOv2 tự giám sát mạnh mẽ, chúng ta có thể giảm đáng kể chi phí tính toán và thời gian huấn luyện. Sử dụng DINOv2, chúng ta chỉ cần thêm một segmentation head (đầu phân vùng) lên trên backbone đã được huấn luyện trước và huấn luyện một vài nghìn tham số để có hiệu suất tốt. Đây chính xác là những gì chúng ta sẽ đề cập trong bài viết này. Chúng ta sẽ sửa đổi backbone DINOv2, thêm một pixel classifier (bộ phân loại pixel) đơn giản lên trên nó và huấn luyện DINOv2 cho phân vùng ngữ nghĩa.
Nhảy đến phần Tải Code
Hình 1: Kết quả huấn luyện mô hình DINOv2 cho phân vùng ngữ nghĩa.
Chúng ta sẽ đề cập đến những gì khi huấn luyện DINOV2 cho phân vùng ngữ nghĩa?
Chúng ta sẽ chủ yếu trả lời các câu hỏi sau:
- Tại sao chúng ta cần bài viết này?
- Làm thế nào để sửa đổi backbone DINOv2 cho phân vùng ngữ nghĩa?
- Quy trình huấn luyện mô hình là gì?
- Chúng ta có thể mong đợi những kết quả nào sau khi huấn luyện DINOv2 cho phân vùng ngữ nghĩa?
- Những hạn chế của phương pháp tiếp cận của chúng ta là gì?
Lưu ý: Đây không phải là code sạch nhất có thể để chuyển đổi backbone DINOv2 thành một mô hình phân vùng ngữ nghĩa. Đây là codebase giai đoạn đầu cho một dự án lớn hơn mà chúng ta sẽ sớm thực hiện tại DebuggerCafe. Code không hoàn toàn sạch, rất nhiều “vá” đã được đưa vào phần mô hình hóa, với một số phần vẫn chưa được trả lời.
Tại sao chúng ta cần bài viết này?
Kho lưu trữ DINOv2 chính thức cung cấp hai notebooks cho các mô hình DINOv2 đã được sửa đổi, một cho phân vùng ngữ nghĩa và một cho ước tính độ sâu. Chúng ta sẽ tập trung hoàn toàn vào phân vùng ngữ nghĩa ở đây.
Kho lưu trữ chỉ cung cấp code suy luận cho phân vùng ngữ nghĩa với tích hợp MMSegmentation. Hiện tại, rất khó để thiết lập kho lưu trữ với các yêu cầu MMSegmentation vì một số vấn đề về dependency. Hơn nữa, không rõ làm thế nào để huấn luyện mô hình cho phân vùng ngữ nghĩa và chọn siêu tham số nào. Ngoài ra, kho lưu trữ cung cấp các segmentation head được huấn luyện trên Pascal VOC và ADE20k, cùng với một mô hình DINOv2 được huấn luyện lại hoàn toàn dựa trên segmentation head Mask2Former.
Để giảm thiểu một số vấn đề trên, bài viết này nhằm mục đích loại bỏ hoàn toàn dependency MMSegmentation. Chúng ta dựa vào code backbone PyTorch thuần túy từ kho lưu trữ DINOv2 trong khi thêm code cần thiết để xây dựng phiên bản phân vùng ngữ nghĩa của nó.
Vì dự án đang trong giai đoạn đầu, code vẫn chưa hoàn toàn sạch như chúng tôi đã đề cập ở trên. Một số siêu tham số và phương pháp tiếp cận vẫn chưa được hoàn thiện. Tuy nhiên, với bài viết này, cộng đồng cũng có được một cái gì đó hữu hình mà họ có thể nghịch, đề xuất và cải thiện.
Bộ Dữ Liệu Phân Vùng Người
Chúng ta sẽ thử nghiệm với một bộ dữ liệu rất đơn giản, phân vùng người nói chính xác, trong bài viết này. Nỗ lực của chúng ta trong việc khám phá DINOv2 cho các tác vụ khác nhau (bắt đầu với phân vùng trong bài viết này) có thể không tạo ra kết quả rất chính xác trong lần thử đầu tiên, nhưng cuối cùng chúng ta sẽ đạt được điều đó.
Chúng ta sẽ huấn luyện mô hình trên bộ dữ liệu Phân Vùng Người Đi Bộ Penn-Fudan. Đây là một bộ dữ liệu nhỏ chỉ với 146 mẫu huấn luyện và 24 mẫu xác thực. Điều này là hoàn hảo để thử nghiệm với một thay đổi code lớn như vậy để chúng ta có thể lặp lại nhanh chóng.
Các mask chứa một mask nhãn khác nhau cho mỗi người vì nó đã được điều chỉnh từ một bộ dữ liệu phân vùng instance (instance segmentation). Trong quá trình huấn luyện, chúng ta sẽ chuyển đổi mask của mỗi người thành giá trị 255 và background thành 0.
Bạn sẽ thấy cấu trúc sau sau khi tải xuống và giải nén dữ liệu từ liên kết trên.
PennFudanPed/
├── train_images
├── train_masks
├── valid_images
└── valid_masks
content_copy download Use code with caution.
Sau đây là một số hình ảnh và mask ground truth (nhãn mask người đi bộ đã được chuyển đổi thành giá trị 255).
Hình 2: Bản đồ phân vùng ground truth từ bộ dữ liệu Người Đi Bộ Penn-Fudan.
Cấu Trúc Thư Mục Dự Án
Hãy xem cấu trúc thư mục dự án trước khi chuyển sang phần code.
├── input
│ ├── inference_data
│ └── PennFudanPed
├── outputs
│ ├── valid_preds
│ ├── best_model_iou.pth
│ └── best_model_loss.pth
├── config.py
├── datasets.py
├── engine.py
├── infer_image.py
├── infer_video.py
├── metrics.py
├── model.py
├── requirements.txt
├── temp.py
├── train.py
└── utils.py
content_copy download Use code with caution.
- Thư mục input chứa bộ dữ liệu phân vùng người đi bộ Penn-Fudan mà chúng ta đã tải xuống ở trên.
- Thư mục outputs sẽ chứa các kết quả huấn luyện, trọng số mô hình đã huấn luyện, cũng như các kết quả suy luận.
- Trong thư mục dự án gốc, chúng ta có tất cả các file Python mà chúng ta sẽ làm việc.
Tất cả các file code, bộ dữ liệu suy luận và trọng số đã huấn luyện sẽ có sẵn thông qua phần “Tải Code”.
Tải Code
Bạn có thể tải source code cho hướng dẫn này [tại đây](Liên kết tải xuống).
Cài Đặt Dependencies
Chúng ta có thể cài đặt tất cả các dependency cần thiết bằng file requirements.txt.
pip install -r requirements.txt
content_copy download Use code with caution.
Nếu bạn gặp lỗi liên quan đến OpenCV (vì Albumentations), vui lòng thực thi lệnh sau sau khi cài đặt các yêu cầu trên.
pip uninstall opencv-python opencv-python-headless
pip install opencv-python
content_copy download Use code with caution.
Phân Vùng Ngữ Nghĩa Sử Dụng DINOv2
Bắt đầu từ phần này, chúng ta sẽ đề cập đến các phần code quan trọng liên quan đến dự án. Như chúng ta đã thảo luận trước đó, rất nhiều thành phần code hiện đang được “vá” một cách nghiêm ngặt và có các phương pháp thực hành code tốt hơn cho chúng. Chúng ta sẽ chỉ ra những điều này trong khi đề cập đến code mô hình.
Làm thế nào để sửa đổi Backbone DINOv2 cho phân vùng ngữ nghĩa?
Code để tạo mô hình phân vùng dựa trên DINOv2 nằm trong file model.py. Hãy bắt đầu với các câu lệnh import.
# model.py
import torch
import urllib
import warnings
import torch.nn.functional as F
import torch.nn as nn
from functools import partial
from collections import OrderedDict
from torchinfo import summary
content_copy download Use code with caution.Python
Xác Định Backbone DINOv2
Bước đầu tiên để tạo mô hình là xác định kiến trúc mô hình và tải xuống backbone đã huấn luyện trước.
# model.py
BACKBONE_SIZE = "small" # in ("small", "base", "large" or "giant")
backbone_archs = {
"small": "vits14",
"base": "vitb14",
"large": "vitl14",
"giant": "vitg14",
}
backbone_arch = backbone_archs[BACKBONE_SIZE]
backbone_name = f"dinov2_{backbone_arch}"
HEAD_SCALE_COUNT = 3 # more scales: slower but better results, in (1,2,3,4,5)
HEAD_DATASET = "voc2012" # in ("ade20k", "voc2012")
HEAD_TYPE = "ms" # in ("ms, "linear")
DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
head_config_url = f"{DINOV2_BASE_URL}/{backbone_name}/{backbone_name}_{HEAD_DATASET}_{HEAD_TYPE}_config.py"
head_checkpoint_url = f"{DINOV2_BASE_URL}/{backbone_name}/{backbone_name}_{HEAD_DATASET}_{HEAD_TYPE}_head.pth"
backbone_model = torch.hub.load(repo_or_dir="facebookresearch/dinov2", model=backbone_name)
backbone_model.cuda()
content_copy download Use code with caution.Python
Chúng ta đang chọn backbone nhỏ từ DINOv2. Mỗi tên backbone được ghép nối với một mô hình ViT khác nhau như chúng ta có thể thấy trong backbone_archs. Đối với mô hình DINOv2-Small, đó là vits14, là một mô hình Vision Transformer nhỏ với các patch có độ phân giải 14×14.
Từ đó, chúng ta xác định backbone_arch và backbone_name, từ đó xác định cấu hình head và URL checkpoint head.
Chúng ta cũng có HEAD_SCALE_COUNT, HEAD_DATASET và HEAD_TYPE. HEAD_DATASET xác định bộ dữ liệu mà segmentation head đã được huấn luyện trên đó. Trong trường hợp này, đó là bộ dữ liệu Pascal VOC. HEAD_TYPE là đa tỷ lệ, cho mô hình biết sau này để tổng hợp các đặc trưng từ các lớp khác nhau trước khi chuyển nó sang lớp phân loại pixel.
Hiện tại, mặc dù chúng ta đã xác định HEAD_SCALE_COUNT, chúng ta không sử dụng nó. Đây là một phần của file cấu hình MMCV và việc sử dụng MMSegmentation cho huấn luyện/suy luận sẽ tự động sử dụng điều này. Tôi vẫn chưa tìm ra cách sử dụng nó mà không có file cấu hình MMSegmentation để huấn luyện hoặc suy luận. Một suy đoán là chúng ta nên có các tỷ lệ hình ảnh khác nhau cho mỗi batch để huấn luyện mô hình ở các độ phân giải khác nhau. Tuy nhiên, cần xem xét kỹ hơn codebase nội bộ để triển khai thích hợp.
Tải Cấu Hình Backbone
Bây giờ đến phần lộn xộn của code. Chúng ta cần tải xuống file cấu hình để xác định đúng backbone mô hình. Các tác giả ban đầu đã sử dụng cấu hình này cho huấn luyện MMSegmentation, đây là phương pháp dự kiến cho thư viện.
# model.py
# TODO: This part needs cleaning #
def load_config_from_url(url: str) -> str:
with urllib.request.urlopen(url) as f:
return f.read().decode()
cfg_str = load_config_from_url(head_config_url)
with open('temp.py', 'w') as f:
f.write(cfg_str)
from temp import model as model_dict
##################################
backbone_model.forward = partial(
backbone_model.get_intermediate_layers,
n=model_dict['backbone']['out_indices'],
reshape=True,
)
for name, param in backbone_model.named_parameters():
param.requires_grad = False
content_copy download Use code with caution.Python
Chúng ta tải xuống cấu hình mô hình bằng head_config_url và lưu nó dưới dạng temp.py. File này chứa một số từ điển xác định các cấu hình huấn luyện, suy luận và thử nghiệm quan trọng cùng với các cấu hình mô hình.
Chúng ta tải từ điển cấu hình mô hình làm model_dict được sử dụng thêm để xác định một phương pháp partial cùng với số lượng lớp (n) để lấy các đặc trưng từ đó. Đối với phân vùng, chúng ta sẽ nối các đặc trưng từ 4 lớp khác nhau được xác định là out_indices trong file cấu hình (temp.py).
Sau đây hiển thị cấu hình mô hình từ temp.py.
model = dict(
type='EncoderDecoder',
pretrained=None,
backbone=dict(type='DinoVisionTransformer', out_indices=[8, 9, 10, 11]),
decode_head=dict(
type='BNHead',
in_channels=[384, 384, 384, 384],
in_index=[0, 1, 2, 3],
input_transform='resize_concat',
channels=1536,
dropout_ratio=0,
num_classes=21,
norm_cfg=dict(type='SyncBN', requires_grad=True),
align_corners=False,
loss_decode=dict(
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
test_cfg=dict(mode='slide', crop_size=(640, 640), stride=(320, 320)))
content_copy download Use code with caution.Python
Như chúng ta có thể thấy, cấu hình hướng dẫn phương pháp partial trích xuất các đặc trưng từ các lớp 8, 9, 10 và 11. Điều quan trọng cần nhớ là cấu hình này chỉ hợp lệ cho mô hình ViT-Small và sẽ khác nhau cho mỗi backbone.
Hàm Thay Đổi Kích Thước Tùy Chỉnh
Tiếp theo, chúng ta có một hàm thay đổi kích thước tùy chỉnh được sử dụng trong segmentation head của mô hình.
# model.py
def resize(input_data,
size=None,
scale_factor=None,
mode='nearest',
align_corners=None,
warning=True):
if warning:
if size is not None and align_corners:
input_h, input_w = tuple(int(x) for x in input.shape[2:])
output_h, output_w = tuple(int(x) for x in size)
if output_h > input_h or output_w > input_w:
if ((output_h > 1 and output_w > 1 and input_h > 1
and input_w > 1) and (output_h - 1) % (input_h - 1)
and (output_w - 1) % (input_w - 1)):
warnings.warn(
f'When align_corners={align_corners}, '
'the output would more aligned if '
f'input size {(input_h, input_w)} is `x+1` and '
f'out size {(output_h, output_w)} is `nx+1`')
return F.interpolate(input_data, size, scale_factor, mode, align_corners)
content_copy download Use code with caution.Python
BNHead cho Phân Loại Pixel
Chúng ta sử dụng lớp BNHead từ kho lưu trữ DINOv2 để nối các đặc trưng đã trích xuất, xác định segmentation head phân loại pixel và thực hiện forward pass.
# model.py
class BNHead(nn.Module):
"""Just a batchnorm."""
def __init__(self, resize_factors=None, **kwargs):
super().__init__(**kwargs)
# HARDCODED IN_CHANNELS FOR NOW.
self.in_channels = 1536
self.bn = nn.SyncBatchNorm(self.in_channels)
self.resize_factors = resize_factors
self.in_index = [0, 1, 2, 3]
self.input_transform = 'resize_concat'
self.align_corners = False
self.conv_seg = nn.Conv2d(self.in_channels, 21, kernel_size=1)
def _forward_feature(self, inputs):
"""Forward function for feature maps before classifying each pixel with
``self.cls_seg`` fc.
Args:
inputs (list[Tensor]): List of multi-level img features.
Returns:
feats (Tensor): A tensor of shape (batch_size, self.channels,
H, W) which is feature map for last layer of decoder head.
"""
x = self._transform_inputs(inputs)
feats = self.bn(x)
return feats
def _transform_inputs(self, inputs):
"""Transform inputs for decoder.
Args:
inputs (list[Tensor]): List of multi-level img features.
Returns:
Tensor: The transformed inputs
"""
if self.input_transform == "resize_concat":
# accept lists (for cls token)
input_list = []
for x in inputs:
if isinstance(x, list):
input_list.extend(x)
else:
input_list.append(x)
inputs = input_list
# an image descriptor can be a local descriptor with resolution 1x1
for i, x in enumerate(inputs):
if len(x.shape) == 2:
inputs[i] = x[:, :, None, None]
# select indices
inputs = [inputs[i] for i in self.in_index]
# Resizing shenanigans
# print("before", *(x.shape for x in inputs))
if self.resize_factors is not None:
assert len(self.resize_factors) == len(inputs), (len(self.resize_factors), len(inputs))
inputs = [
resize(input=x, scale_factor=f, mode="bilinear" if f >= 1 else "area")
for x, f in zip(inputs, self.resize_factors)
]
# print("after", *(x.shape for x in inputs))
upsampled_inputs = [
resize(input_data=x, size=inputs[0].shape[2:], mode="bilinear", align_corners=self.align_corners)
for x in inputs
]
inputs = torch.cat(upsampled_inputs, dim=1)
elif self.input_transform == "multiple_select":
inputs = [inputs[i] for i in self.in_index]
else:
inputs = inputs[self.in_index]
return inputs
def cls_seg(self, feat):
"""Classify each pixel."""
output = self.conv_seg(feat)
return output
def forward(self, inputs):
"""Forward function."""
output = self._forward_feature(inputs)
output = self.cls_seg(output)
return output
content_copy download Use code with caution.Python
Lớp phân loại pixel chỉ là một lớp Conv2D với 21 lớp (cho VOC). Chúng ta sẽ thay đổi các kênh đầu ra theo số lượng lớp sau. Phương thức forward biến đổi các đầu vào theo đầu ra đa tỷ lệ, thay đổi kích thước chúng và chuyển chúng qua segmentation head.
Một yếu tố cần lưu ý ở đây là in_channels cho segmentation head mà chúng ta hardcode là 1536. Đây là số lượng đặc trưng đầu ra cuối cùng sau khi nối các đặc trưng có 384 chiều 4 lần (vì chúng ta nối các đặc trưng từ 4 lớp). Sau này, sẽ tốt hơn nhiều nếu làm cho tính toán này trở nên động.
Xác Định Mô Hình Phân Vùng DINOv2 Cuối Cùng
Cuối cùng, chúng ta xác định mô hình Phân Vùng DINOv2 dưới dạng một từ điển Sequential Ordered và Named với một backbone và decode_head.
# model.py
class DINOv2Segmentation(nn.Module):
def __init__(self):
super(DINOv2Segmentation, self).__init__()
self.backbone_model = backbone_model
self.decode_head = BNHead()
self.model = nn.Sequential(OrderedDict([
('backbone', self.backbone_model),
('decode_head', self.decode_head)
]))
def forward(self, x):
outputs = self.model(x)
return outputs
if __name__ == '__main__':
model = DINOv2Segmentation()
summary(
model,
(1, 3, 644, 644),
col_names=('input_size', 'output_size', 'num_params'),
row_settings=['var_names']
)
content_copy download Use code with caution.Python
Thực thi file này với lệnh python model.py sẽ cho kết quả sau.
=============================================================================================================================
Layer (type (var_name)) Input Shape Output Shape Param #
=============================================================================================================================
DINOv2Segmentation (DINOv2Segmentation) [1, 3, 644, 644] [1, 21, 46, 46] --
├─Sequential (model) [1, 3, 644, 644] [1, 21, 46, 46] --
│ └─DinoVisionTransformer (backbone) [1, 3, 644, 644] [1, 384, 46, 46] 526,848
│ │ └─PatchEmbed (patch_embed) [1, 3, 644, 644] [1, 2116, 384] (226,176)
│ │ └─ModuleList (blocks) -- -- (21,302,784)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (768)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (recursive)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (recursive)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (recursive)
│ └─BNHead (decode_head) [1, 384, 46, 46] [1, 21, 46, 46] --
│ │ └─SyncBatchNorm (bn) [1, 1536, 46, 46] [1, 1536, 46, 46] 3,072
│ │ └─Conv2d (conv_seg) [1, 1536, 46, 46] [1, 21, 46, 46] 32,277
=============================================================================================================================
Total params: 22,091,925
Trainable params: 35,349
Non-trainable params: 22,056,576
Total mult-adds (Units.MEGABYTES): 568.19
=============================================================================================================================
Input size (MB): 4.98
Forward/backward pass size (MB): 1073.41
Params size (MB): 86.26
Estimated Total Size (MB): 1164.64
=============================================================================================================================
content_copy download Use code with caution.
Vì chúng ta đang đóng băng backbone, chỉ có decode head là có thể huấn luyện được với 35.349 tham số cho 21 lớp. Điều này sẽ giảm hơn nữa khi chúng ta chỉ huấn luyện một lớp, đó là người.
Một điểm cuối cùng cần lưu ý là kích thước hình ảnh 644×644 mà chúng ta đã sử dụng trong dummy forward pass ở trên. Do cấu trúc của backbone (kích thước kernel và padding), các đặc trưng đầu vào bị cắt và độ phân giải 640×640 gây ra lỗi vì các patch được tạo ra không chia hết cho 14. Tuy nhiên, trong quá trình huấn luyện, chúng ta sẽ xử lý điều này trong dataloader image transforms để chúng ta có thể truyền độ phân giải 640×640 trong các đối số dòng lệnh để tránh nhầm lẫn.
Các Yếu Tố Quan Trọng Về Chuẩn Bị Bộ Dữ Liệu
Trong phần này, chúng ta sẽ khám phá một số đoạn code quan trọng từ file datasets.py chuẩn bị bộ dữ liệu.
Đầu tiên là các giá trị chuẩn hóa. Chúng ta không có các giá trị chuẩn hóa cho bộ dữ liệu LVD-142M. Vì vậy, chúng ta sử dụng các giá trị chuẩn hóa ImageNet.
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
content_copy download Use code with caution.Python
Thứ hai là các transform huấn luyện và xác thực.
# TODO: Batchwise rescaling with different image ratios.
def train_transforms(img_size):
"""
Transforms/augmentations for training images and masks.
:param img_size: Integer, for image resize.
"""
train_image_transform = A.Compose([
A.Resize(img_size[1], img_size[0], always_apply=True),
A.PadIfNeeded(
min_height=img_size[1]+4,
min_width=img_size[0]+4,
position='center',
value=0,
mask_value=0
),
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2),
A.Rotate(limit=25),
A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.)
], is_check_shapes=False)
return train_image_transform
# TODO: Batchwise rescaling with different image ratios.
def valid_transforms(img_size):
"""
Transforms/augmentations for validation images and masks.
:param img_size: Integer, for image resize.
"""
valid_image_transform = A.Compose([
A.Resize(img_size[1], img_size[0], always_apply=True),
A.PadIfNeeded(
min_height=img_size[1]+4,
min_width=img_size[0]+4,
position='center',
value=0,
mask_value=0
),
A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.)
], is_check_shapes=False)
return valid_image_transform
content_copy download Use code with caution.Python
Sau đây là các transform phổ biến cho huấn luyện và xác thực:
- Đầu tiên, chúng ta thay đổi kích thước hình ảnh thành các giá trị thay đổi kích thước được truyền qua dòng lệnh trong quá trình thực thi script huấn luyện. Đối với dự án này, nó sẽ là 640×640.
- Sau đó, chúng ta căn giữa mỗi hình ảnh và mask bằng 4 pixel đen. Điều này mang lại độ phân giải cuối cùng là 644×644. Việc triển khai ban đầu từ các tác giả thực hiện điều này một cách nhanh chóng và là một phần của pipeline mô hình. Hơn nữa, các tác giả triển khai padding động để mô hình không bao giờ gây ra lỗi chia patch. Trong trường hợp của chúng ta, việc căn giữa 4 pixel sẽ hoạt động cho hình ảnh 640×640. Nó được hardcode cho bây giờ. Nói chung, nó cần phải là image_size + (image_size/160). Vì vậy, đối với kích thước hình ảnh 1280×1280, chúng ta nên có padding tâm là 8 pixel.
- Chúng ta áp dụng lật ngang, độ sáng tương phản ngẫu nhiên và các augmentation xoay cho bộ huấn luyện để tránh overfitting.
Hơn nữa, file cũng chứa lớp SegmentationDataset để chuẩn bị các bộ dữ liệu tùy chỉnh. Vì tất cả các instance của người có các giá trị khác nhau, chúng ta sẽ tạo mỗi pixel có giá trị lớn hơn 0 là 255. Đoạn code này trong phương thức __getitem__ của lớp bộ dữ liệu sẽ thực hiện công việc này.
# Make all pixel > 0 as 255.
im = mask > 0
mask[im] = 255
mask[np.logical_not(im)] = 0
content_copy download Use code with caution.Python
Logits Đến Logic Bản Đồ Pixel Được Upsample
Mô hình ban đầu trả về các logits có hình dạng 46×46 dưới dạng bản đồ pixel thô. Chúng được downsample 14 lần so với đầu vào 644×644 mà chúng ta truyền. Để tính toán loss và cũng để trực quan hóa suy luận, chúng ta cần các logits được upsample sẽ được chuyển đổi thành bản đồ phân vùng.
Để đạt được điều này, trong các phương thức train và validate của engine.py, chúng ta có logic sau.
upsampled_logits = nn.functional.interpolate(
outputs, size=target.shape[-2:],
mode="bilinear",
align_corners=False
)
content_copy download Use code with caution.Python
Chúng ta áp dụng nội suy song tuyến tính (bilinear interpolation) để thay đổi kích thước các logits thành hình dạng hình ảnh đầu vào ban đầu. Logic tương tự được sử dụng trong quá trình suy luận.
Điều này hoàn thành tất cả các đoạn code quan trọng mà chúng ta cần đề cập.
Huấn Luyện Mô Hình DINOv2 cho Phân Vùng Ngữ Nghĩa
Trong script train.py, chúng ta cần sửa đổi số lượng lớp cho segmentation head cuối cùng sau khi khởi tạo mô hình.
model = DINOv2Segmentation()
model.decode_head.conv_seg = nn.Conv2d(1536, len(ALL_CLASSES), kernel_size=(1, 1), stride=(1, 1))
content_copy download Use code with caution.Python
Chúng ta đã sẵn sàng để bắt đầu quá trình huấn luyện. Thực thi lệnh sau trong terminal trong thư mục dự án gốc.
python train.py --lr 0.001 --batch 20 --imgsz 640 640 --epochs 65 --scheduler --scheduler-epochs 50
content_copy download Use code with caution.
Chúng ta sử dụng learning rate cao hơn bình thường vì chúng ta chỉ huấn luyện segmentation head mới và các trọng số backbone bị đóng băng. Hơn nữa, điều này tuân theo cùng một pipeline như file cấu hình ban đầu. Chúng ta huấn luyện với hình ảnh có độ phân giải 640×640 trong 65 epoch và áp dụng một learning rate scheduler để giảm learning rate đi 10 lần sau 50 epoch.
Sau đây là các kết quả được cắt bớt.
=============================================================================================================================
Layer (type (var_name)) Input Shape Output Shape Param #
=============================================================================================================================
DINOv2Segmentation (DINOv2Segmentation) [1, 3, 644, 644] [1, 2, 46, 46] --
├─Sequential (model) [1, 3, 644, 644] [1, 2, 46, 46] --
│ └─DinoVisionTransformer (backbone) [1, 3, 644, 644] [1, 384, 46, 46] 526,848
│ │ └─PatchEmbed (patch_embed) [1, 3, 644, 644] [1, 2116, 384] (226,176)
│ │ └─ModuleList (blocks) -- -- (21,302,784)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (768)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (recursive)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (recursive)
│ │ └─LayerNorm (norm) [1, 2117, 384] [1, 2117, 384] (recursive)
│ └─BNHead (decode_head) [1, 384, 46, 46] [1, 2, 46, 46] --
│ │ └─SyncBatchNorm (bn) [1, 1536, 46, 46] [1, 1536, 46, 46] 3,072
│ │ └─Conv2d (conv_seg) [1, 1536, 46, 46] [1, 2, 46, 46] 3,074
=============================================================================================================================
Total params: 22,062,722
Trainable params: 6,146
Non-trainable params: 22,056,576
Total mult-adds (Units.MEGABYTES): 506.39
=============================================================================================================================
Input size (MB): 4.98
Forward/backward pass size (MB): 1073.08
Params size (MB): 86.14
Estimated Total Size (MB): 1164.20
=============================================================================================================================
EPOCH: 1
Training
100%|████████████████████| 8/8 [00:10<00:00, 1.34s/it]
Validating
100%|████████████████████| 2/2 [00:02<00:00, 1.15s/it]
content_copy download Use code with caution.
