بخش‌بندی معنایی چند کلاسه با استفاده از DINOv2

اگرچه DINOv2 بک‌بون‌های از پیش آموزش‌دیده قدرتمندی را ارائه می‌دهد، آموزش آن برای خوب بودن در وظایف بخش‌بندی معنایی می‌تواند دشوار باشد. فقط آموزش یک سر بخش‌بندی ممکن است در مواقعی نتایج نامطلوبی به همراه داشته باشد. در این مقاله، ما بر دو نکته تمرکز خواهیم کرد: بخش‌بندی معنایی چند کلاسه با استفاده از DINOv2 و مقایسه نتایج با آموزش فقط بخش‌بندی و تنظیم دقیق کل شبکه.

پرش به بخش دانلود کد

در مقالات قبلی در مورد DINOv2 ، ما دو جنبه را پوشش داده‌ایم:

• طبقه‌بندی تصویر با استفاده از DINOv2 که در آن نتایج بین تنظیم دقیق و یادگیری انتقالی را مقایسه می‌کنیم.
• آموزش یک سر بخش‌بندی بالای ویژگی‌های DINOv2 برای بخش‌بندی افراد.

در این مقاله، یک گام فراتر خواهیم رفت و مدل را روی یک مجموعه داده بخش‌بندی چند کلاسه تنظیم دقیق خواهیم کرد. در این راستا، دو آزمایش انجام خواهیم داد و نتایج را مقایسه خواهیم کرد.

• ما با بحث در مورد مجموعه داده شروع خواهیم کرد. در اینجا از مجموعه داده بخش‌بندی Pascal VOC استفاده خواهیم کرد.
• بعد، کد آموزشی را تنظیم خواهیم کرد. بخش‌های کد اصلاح‌شده را با مقاله بخش‌بندی قبلی در مورد DINOv2 مقایسه خواهیم کرد.
• دو آزمایش را اجرا خواهیم کرد:
  • آموزش فقط سر بخش‌بندی.
  • تنظیم دقیق کل شبکه.
• در نهایت، نتایج را تجزیه و تحلیل خواهیم کرد.

ما آزمایش‌های آموزشی خود را برای بخش‌بندی چند کلاسه DINOv2 بر روی مجموعه داده بخش‌بندی Pascal VOC اجرا خواهیم کرد.

می‌توانید مجموعه داده را از اینجا در Kaggle دانلود کنید.

در زیر ساختار دایرکتوری مجموعه داده پس از دانلود و استخراج مجموعه داده آمده است.

voc_2012_segmentation_data/
+-- train_images
+-- train_labels
+-- valid_images
+-- valid_labels

این مجموعه داده شامل ۱۴۶۴ تصویر آموزشی و جفت ماسک و ۱۴۴۹ تصویر اعتبارسنجی و جفت ماسک است.

در اینجا چند نمونه آورده شده است.

این مجموعه داده شامل ۲۱ کلاس از جمله کلاس پس‌زمینه است.

[
'background',
'aeroplane',
'bicycle',
'bird',
'boat',
'bottle',
'bus',
'car',
'cat',
'chair',
'cow',
'dining table',
'dog',
'horse',
'motorbike',
'person',
'potted plant',
'sheep',
'sofa',
'train',
'tv/monitor'
]

برای هر کلاس، یک نگاشت بخش‌بندی رنگ RGB داریم.

[
[0, 0, 0],
[128, 0, 0],
[0, 128, 0],
[128, 128, 0],
[0, 0, 128],
[128, 0, 128],
[0, 128, 128],
[128, 128, 128],
[64, 0, 0],
[192, 0, 0],
[64, 128, 0],
[192, 128, 0],
[64, 0, 128],
[192, 0, 128],
[64, 128, 128],
[192, 128, 128],
[0, 64, 0],
[128, 64, 0],
[0, 192, 0],
[128, 192, 0],
[0, 64, 128]
]

به دست آوردن نتایج خوب با آموزش مستقیم یک مدل بخش‌بندی فقط با استفاده از یک بک‌بون از پیش آموزش‌دیده ImageNet در این مجموعه داده دشوار است. این به این دلیل است که اگرچه مجموعه داده حاوی تنوع است، اما تعداد نمونه‌ها کم است.

علاوه بر این، پس از انجام آزمایش‌های آموزشی، تجزیه و تحلیل خواهیم کرد که چگونه آموزش فقط سر بخش‌بندی و تنظیم دقیق کل شبکه بخش‌بندی DINOv2 بر نتایج تأثیر می‌گذارد.

بیایید نگاهی به ساختار دایرکتوری پروژه بیندازیم.

+-- input
¦ +-- inference_data
¦ +-- voc_2012_segmentation_data
+-- notebooks
¦ +-- visualize.ipynb
+-- outputs
¦ +-- fine_tuning
¦ +-- inference_results_video
¦ +-- transfer_learning
+-- config.py
+-- datasets.py
+-- engine.py
+-- infer_image.py
+-- infer_video.py
+-- metrics.py
+-- model_config.py
+-- model.py
+-- requirements.txt
+-- train.py
+-- utils.py

• دایرکتوری input شامل داده‌های آموزش و استنتاج است.
• دایرکتوری outputs شامل نتایج پس از آموزش مدل و همچنین نتایج استنتاج است.
• در دایرکتوری ریشه پروژه، همه فایل‌های پایتون مورد نیاز برای آموزش مدل بخش‌بندی DINOv2 را داریم. در میان این موارد، ما در اینجا بر روی فایل model.py تمرکز خواهیم کرد.

مدل‌های آموزش‌دیده، فایل‌های کد و داده‌های استنتاج از طریق بخش دانلود در دسترس هستند. برای آموزش مدل، لطفاً مجموعه داده را از Kaggle دانلود کرده و در ساختار دایرکتوری بالا مرتب کنید.

دانلود کد منبع برای این آموزش

پس از دانلود فایل کد و استخراج آن، می‌توانید وابستگی‌ها را با استفاده از فایل requirements نصب کنید.

pip install -r requirements.txt

اکنون، بیایید به بخش کدنویسی مقاله بپردازیم.

ما در این مقاله عمدتاً به کد ساخت مدل خواهیم پرداخت. بیشتر کدها مشابه مواردی است که در یکی از مقالات قبلی خود در مورد آزمایش‌های بخش‌بندی معنایی DINOv2 با یادگیری انتقالی و تنظیم دقیق بحث کردیم. این مقاله برخی از اجزای ساخت مدل را در مقایسه با اولین مقاله بخش‌بندی DINOv2 که قبلاً ذکر شد، ساده کرده است.

کد ساخت مدل در فایل model.py قرار دارد.

در زیر واردات مورد نیاز ما آمده است.

import torch
import torch.nn as nn

from functools import partial
from collections import OrderedDict
from torchinfo import summary
from model_config import model as model_dict

ما پیکربندی‌های مدل را از ماژول model_config وارد می‌کنیم. این شامل تمام پیکربندی‌های مدلی است که بخشی از مخزن اصلی DINOv2 بودند که در آن از کتابخانه MMSegmentation استفاده شده بود. رویکرد ما این فرآیند را ساده می‌کند و نیاز MMSegmentation را به طور کامل حذف می‌کند.

بعد، یک تابع کمکی برای بارگیری بک‌بون DINOv2 داریم.

def load_backbone(backbone_size="small"):
    backbone_archs = {
        "small": "vits14",
        "base": "vitb14",
        "large": "vitl14",
        "giant": "vitg14",
    }

    backbone_arch = backbone_archs[backbone_size]
    backbone_name = f"dinov2_{backbone_arch}"


    backbone_model = torch.hub.load(repo_or_dir="facebookresearch/dinov2", model=backbone_name)
    backbone_model.cuda()

    backbone_model.forward = partial(
        backbone_model.get_intermediate_layers,
        n=model_dict['backbone']['out_indices'],
        reshape=True,
    )

    return backbone_model

ما از بک‌بون کوچک DINOv2 در اینجا برای آموزش سریع‌تر استفاده می‌کنیم. این شامل حدود ۲۲ میلیون پارامتر است.

ما از یک سر بخش‌بندی خطی ساده استفاده می‌کنیم که در آن توکن‌های خروجی را برای تغذیه به یک لایه کانولوشن نهایی تغییر شکل می‌دهیم.

class LinearClassifierToken(torch.nn.Module):
    def __init__(self, in_channels, nc=1, tokenW=32, tokenH=32):
        super(LinearClassifierToken, self).__init__()
        self.in_channels = in_channels
        self.W = tokenW
        self.H = tokenH
        self.nc = nc
        self.conv = torch.nn.Conv2d(in_channels, nc, (1, 1))

    def forward(self,x):
        outputs =  self.conv(
            x.reshape(-1, self.in_channels, self.H, self.W)
        )
        return outputs

سپس کلاس ساخت مدل نهایی را داریم که همه اجزا را ترکیب می‌کند.

class Dinov2Segmentation(nn.Module):
    def __init__(self, fine_tune=False):
        super(Dinov2Segmentation, self).__init__()

        self.backbone_model = load_backbone()
        print(fine_tune)
        if fine_tune:
            for name, param in self.backbone_model.named_parameters():
                param.requires_grad = True
        else:
            for name, param in self.backbone_model.named_parameters():
                param.requires_grad = False

        self.decode_head = LinearClassifierToken(in_channels=1536, nc=21, tokenW=46, tokenH=46)

        self.model = nn.Sequential(OrderedDict([
            ('backbone', self.backbone_model),
            ('decode_head', self.decode_head)
        ]))

    def forward(self, x):
        features = self.model.backbone(x)

        # `features` is a tuple.
        concatenated_features = torch.cat(features, 1)

        classifier_out = self.decode_head(concatenated_features)

        return classifier_out
    
if __name__ == '__main__':
    model = Dinov2Segmentation()

    summary(
        model, 
        (1, 3, 644, 644),
        col_names=('input_size', 'output_size', 'num_params'),
        row_settings=['var_names']
    )

ما تعداد کلاس‌ها را در اینجا به ۲۱ کدگذاری کرده‌ایم که با مجموعه داده Pascal VOC مطابقت دارد. با این حال، بهتر است این را به عنوان یک آرگومان هنگام ساخت مدل منتقل کنیم.

ما در مقالات قبلی خود در مورد ظرایف ساخت بخش‌بندی DINOv2 بحث کرده‌ایم. در صورت نیاز حتماً به آنها نگاهی بیندازید.

ما می‌توانیم فایل را با استفاده از دستور زیر اجرا کنیم که خلاصه شبکه و پارامترها را به ما نشان می‌دهد.

python model.py

=============================================================================================================================
Layer (type (var_name)) Input Shape Output Shape Param #
=============================================================================================================================
Dinov2Segmentation (Dinov2Segmentation) [1, 3, 644, 644] [1, 21, 46, 46] --
+-Sequential (model) -- -- --
¦ +-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)
¦ +-LinearClassifierToken (decode_head) [1, 1536, 46, 46] [1, 21, 46, 46] --
¦ ¦ +-Conv2d (conv) [1, 1536, 46, 46] [1, 21, 46, 46] 32,277
=============================================================================================================================
Total params: 22,088,853
Trainable params: 32,277
Non-trainable params: 22,056,576
Total mult-adds (Units.MEGABYTES): 568.18
=============================================================================================================================
Input size (MB): 4.98
Forward/backward pass size (MB): 1047.40
Params size (MB): 86.25
Estimated Total Size (MB): 1138.63
=============================================================================================================================

کل شبکه عصبی شامل حدود 22M پارامتر است و سر بخش‌بندی شامل 32277 پارامتر است.

ما از افزایش‌های زیر برای مجموعه داده آموزشی استفاده می‌کنیم:

• چرخش افقی
• کنتراست روشنایی تصادفی
• چرخش

ما مدل را با استفاده از بهینه‌ساز AdamW آموزش خواهیم داد و از Cross Entropy به عنوان تابع زیان استفاده خواهیم کرد.

ما با آزمایش یادگیری انتقالی شروع خواهیم کرد.

python train.py --epochs 20 --imgsz 640 640 --out-dir transfer_learning --batch 2

ما مدل را برای 20 دوره با اندازه تصویر 640×640 و اندازه دسته‌ای 2 آموزش می‌دهیم. ما یک اندازه دسته کوچک را انتخاب می‌کنیم تا بتوانیم همان اندازه دسته را برای تنظیم دقیق نیز برای نتایج قابل مقایسه نگه داریم.

در زیر نتایجی که به دست می‌آوریم آمده است.

EPOCH: 1
Training
100%|¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦| 366/366 [00:57<00:00, 6.36it/s]
Validating
100%|¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦| 363/363 [00:51<00:00, 7.08it/s]

Best validation loss: 0.44906933225958146

Saving best model for epoch: 1


Best validation IoU: 0.14743237744729587

Saving best model for epoch: 1

Train Epoch Loss: 1.0304, Train Epoch PixAcc: 0.7609, Train Epoch mIOU: 0.080638
Valid Epoch Loss: 0.4491, Valid Epoch PixAcc: 0.8659 Valid Epoch mIOU: 0.147432
LR for next epoch: [0.0001]
.
.
.
EPOCH: 19
Training
100%|¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦| 366/366 [00:57<00:00, 6.39it/s]
Validating
100%|¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦¦| 363/363 [00:49<00:00, 7.32it/s]

Best validation loss: 0.17451634874808558

Saving best model for epoch: 19


Best validation
IoU: 0.3335922740670536

Saving best model for epoch: 19

Train Epoch Loss: 0.1663, Train Epoch PixAcc: 0.9578, Train Epoch mIOU: 0.272830
Valid Epoch Loss: 0.1745, Valid Epoch PixAcc: 0.9524 Valid Epoch mIOU: 0.333592
LR for next epoch: [0.0001]
EPOCH: 20
Training
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 366/366 [00:57<00:00, 6.36it/s]
Validating
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 363/363 [00:50<00:00, 7.24it/s]

Best validation loss: 0.17389627499521222

Saving best model for epoch: 20


Best validation IoU: 0.34098337358847476

Saving best model for epoch: 20

Train Epoch Loss: 0.1606, Train Epoch PixAcc: 0.9593, Train Epoch mIOU: 0.277356
Valid Epoch Loss: 0.1739, Valid Epoch PixAcc: 0.9526 Valid Epoch mIOU: 0.340983
LR for next epoch: [0.0001]
TESTING
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 363/363 [00:49<00:00, 7.30it/s]
[INFO] Test Loss: 0.17389627499521222
[INFO] Test Pixel Accuracy: 0.9526155591011047
[INFO] Test mIOU: 0.34098337358847476

در اینجا، از آنجا که ما فقط سر بخش‌بندی را آموزش می‌دهیم، یک تفاوت بزرگ بین پارامترهای قابل آموزش و کل پارامترهای شبکه وجود دارد. این به این دلیل است که ما پارامترهای DINOv2 را در اینجا منجمد نگه می‌داریم و بنابراین، مشتقات از طریق آن‌ها عبور نمی‌کنند.

علاوه بر این، ما به دقت پیکسل اعتبارسنجی 95٪ و میانگین IoU (mIOU) حدود 34٪ رسیدیم. mIOU می‌تواند یک معیار گمراه‌کننده باشد. بسته به توزیع داده‌ها، ممکن است خیلی خوب نباشد.

اکنون، بیایید شبکه DINOv2 را تنظیم دقیق کنیم. توجه داشته باشید که تنظیم دقیق شبکه DINOv2 به یک پردازنده گرافیکی قدرتمند نیاز دارد. یک پردازنده گرافیکی با حداقل 24 گیگابایت حافظه VRAM مورد نیاز است.

python train.py --epochs 20 --imgsz 640 640 --out-dir fine_tuning --batch 2 --fine-tune

ما از همان ابرپارامترها به جز استفاده از --fine-tune استفاده خواهیم کرد که به مدل می‌گوید که همه پارامترها را از طریق کد با استفاده از requires_grad=True باز کند. سپس، همه پارامترها در طول آموزش قابل آموزش خواهند بود.

در زیر خروجی برای اجرای فوق آمده است.

EPOCH: 1
Training
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 366/366 [01:31<00:00, 4.01it/s]
Validating
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 363/363 [01:25<00:00, 4.25it/s]

Best validation loss: 0.1744222550732538

Saving best model for epoch: 1


Best validation IoU: 0.3497928817506994

Saving best model for epoch: 1

Train Epoch Loss: 0.2409, Train Epoch PixAcc: 0.9352, Train Epoch mIOU: 0.230641
Valid Epoch Loss: 0.1744, Valid Epoch PixAcc: 0.9509 Valid Epoch mIOU: 0.349793
LR for next epoch: [0.0001]
.
.
.
EPOCH: 19
Training
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 366/366 [01:30<00:00, 4.05it/s]
Validating
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 363/363 [01:25<00:00, 4.26it/s]

Best validation loss: 0.13478370478586374

Saving best model for epoch: 19


Best validation IoU: 0.4343982596498352

Saving best model for epoch: 19

Train Epoch Loss: 0.0705, Train Epoch PixAcc: 0.9787, Train Epoch mIOU: 0.473336
Valid Epoch Loss: 0.1348, Valid Epoch PixAcc: 0.9625 Valid Epoch mIOU: 0.434398
LR for next epoch: [0.0001]
EPOCH: 20
Training
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 366/366 [01:30<00:00, 4.05it/s]
Validating
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 363/363 [01:25<00:00, 4.26it/s]

Best validation loss: 0.13262782954189405

Saving best model for epoch: 20


Best validation IoU: 0.4458368738522375

Saving best model for epoch: 20

Train Epoch Loss: 0.0674, Train Epoch PixAcc: 0.9799, Train Epoch mIOU: 0.482284
Valid Epoch Loss: 0.1326, Valid Epoch PixAcc: 0.9633 Valid Epoch mIOU: 0.445837
LR for next epoch: [0.0001]
TESTING
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 363/363 [01:25<00:00, 4.25it/s]
[INFO] Test Loss: 0.13262782954189405
[INFO] Test Pixel Accuracy: 0.9633225202560425
[INFO] Test mIOU: 0.4458368738522375

ما اکنون دقت پیکسل 96٪ و mIOU 44.5٪ را دریافت می‌کنیم.

آموزش برای هر دوره نیز در اینجا زمان بیشتری می‌برد، زیرا ما اکنون مشتقات را از طریق کل شبکه پخش می‌کنیم.

اکنون، بیایید نگاهی به نتایج استنتاج در تصاویر تست بیندازیم.

برای این کار، می‌توانیم از اسکریپت infer_image.py استفاده کنیم که در آزمایش‌های یادگیری انتقالی و تنظیم دقیق بخش‌بندی DINOv2 توضیح داده‌ایم.

اسکریپت موجود در لینک بالا را به روز کرده‌ایم تا نام کلاس‌ها را با همپوشانی در تصویر نشان دهد. از این رو، می‌توانیم بفهمیم کدام منطقه متعلق به چه کلاسی است. با این حال، در آموزش موجود در این مقاله، ما فقط در حال آموزش یک سر قطعه‌بندی هستیم. از این رو، هیچ همپوشانی کلاس در تصویر مشاهده نخواهید کرد.

شکل‌های ۳ و ۴ نتایج از مدل یادگیری انتقالی و مدل تنظیم دقیق هستند. اگرچه اینها فقط استنتاجات یک‌باره هستند، اما اگر به دقت نگاه کنید، متوجه خواهید شد که تنظیم دقیق در مقایسه با یادگیری انتقالی در تصویر به خوبی بخش‌بندی می‌شود. این بدان معناست که پیکسل‌های قطعه‌بندی شده در بخش‌بندی یادگیری انتقالی در حال مخلوط شدن و ادغام شدن با پیکسل‌های همسایه هستند. در قطعه‌بندی که در شکل ۴ نشان داده شده است، این موارد به‌طور برجسته‌ای کمتر است. به‌طور کلی، این نشان می‌دهد که قطعه‌بندی حاصل از تنظیم دقیق در این مقاله، در مقایسه با قطعه‌بندی به‌دست آمده از یادگیری انتقالی برتر است.

در این مقاله، ما به این سؤال پاسخ دادیم که "آیا تنظیم دقیق DINOv2 به افزایش عملکرد قطعه‌بندی معنایی چند کلاسه کمک می‌کند؟"، یعنی در قطعه‌بندی در اینجا ما تلاشی برای ارائه استدلال و نتایج ملموس و ملموس با به‌کارگیری آن در مجموعه داده قطعه‌بندی Pascal VOC انجام می‌دهیم. از آزمایش‌هایی که انجام دادیم، این‌طور به نظر می‌رسد که قطعه‌بندی دقیق در DINOv2 به این دلیل نتایج بهتری را برای قطعه‌بندی ارائه می‌دهد که همه لایه‌های شبکه عصبی برای ایجاد یک نقشه ویژگی بهینه به‌روز می‌شوند.

می‌توانید این آموزش را به عنوان قطعه‌بندی روی فیلم‌ها، یا افزایش حجم یا ابعاد تصویر به عنوان بخشی از آزمایش، یا گسترش آن به مجموعه داده‌ای متفاوت گسترش دهید.

آنچه در این مقاله متوجه شدیم، این بود که آموزش کل شبکه زمان آموزش را به تأخیر می‌اندازد و نیاز به پردازنده گرافیکی با VRAM بالایی دارد. از این رو، یک تعادل بین هزینه و دقت برای این نوع قطعه‌بندی با استفاده از DINOv2، به عنوان یک مدل یادگیری خود نظارتی وجود دارد.