Можно ли обучить модель обнаружения дефектов кранов без маркированных данных? Отчет о проекте: сравнительное обучение + небольшая доработка

Главным препятствием в обнаружении дефектов деталей подвесных кранов является не разработка модели, а аннотирование данных: съемка 5 000 неаннотированных изображений занимает всего 1–2 дня, тогда как аннотирование аналогичного объема данных требует 2–3 недели. Контрастивное обучение (Contrastive Learning) позволяет выучить универсальные визуальные представления путем построения задачи самообучающегося предварительного обучения на основе ”пар положительных и отрицательных образцов” из неаннотированных данных, после чего для развертывания требуется лишь микронастройка на 100–500 аннотированных изображений. В данной статье на примере двух основных методов контрастивного обучения — SimCLR (2020, Google) и MoCo v2 (2020, FAIR) — проверяется эффективность самообучающегося предварительного обучения с последующей тонкой настройкой на небольшом количестве аннотированных образцов на наборе данных о дефектах деталей грузовых автомобилей, а также проводится сравнение с предварительным обучением на ImageNet и обучением «с нуля». Экспериментальная среда: PyTorch 2.1.0 + NVIDIA A10 (48 ГБ) + 4×RTX 4090 (для работы с большими партиями данных в SimCLR).

Принцип сопоставительного обучения

Основная идея контрастного обучения заключается в том, что ”различные версии одного и того же изображения, подвергнутые обработке данных, должны находиться близко друг к другу в пространстве признаков, а разные изображения — далеко друг от друга”. Возьмем в качестве примера SimCLR: ① К каждому изображению в партии случайным образом применяются два вида обработки данных (случайное кадрирование и масштабирование в диапазоне 0,08–1,0 + цветовое дрожание + гауссово размытие + преобразование в оттенки серого), в результате чего получается 2N обработанных изображений; ② С помощью кодировщика ResNet-18 извлекаются векторы признаков (128-мерные); ③ вычисляется функция потерь NT-Xent (нормализованная температурная шкала перекрестной энтропии): L = -log(exp(sim(z_i, z_j)/τ)/Σexp(sim(z_i, z_k)/τ)), где sim() — косинусная схожесть, τ = 0,5 — температурный коэффициент (чем меньше τ, тем строже различаются отрицательные образцы). Усовершенствования MoCo v2: использование кодировщика с импульсом (momentum = 0,999) и очереди (Queue = 4 096) вместо большого пакета данных SimCLR для решения проблемы недостаточного размера пакета.

Экспериментальные данные и методы оценки

Набор данных: изображения поверхностных дефектов 5 типов компонентов подъемных кранов (разрыв проволок в стальном тросе / износ протектора колес / трещины в рельсах / деформация крюка / растрескивание тормозных накладок + нормальное состояние, всего 6 категорий), всего 1 875 аннотированных изображений (для тонкой настройки и оценки), а также собрано 5 000 неаннотированных изображений для контрастного обучения и предварительной подготовки. В каждой категории примерно по 312 изображений. Эксперименты по тонкой настройке разделены на 4 группы по количеству образцов: 50/100/300/500 изображений на категорию. Показатели оценки: точность Top-1, F1-показатель (макросреднее). Базовая модель: ResNet-18 (при микронастройке всей модели lr = 0,0001, при замораживании FC lr = 0,001).

Основные выводы: ① В сценарии с 50 аннотациями контрастное обучение (F1 = 68,5–70,11 TP3T) значительно превосходит результаты ImageNet (64,81 TP3T) и обучения с нуля (41,81 TP3T); ② Результаты контрастного обучения с 100 аннотированными изображениями (F1 = 89,2–90,31; TP = 3T) приблизительно равны результатам ImageNet с 300 аннотированными изображениями (F1 = 91,11; TP = 3T), что свидетельствует о трёхкратном повышении эффективности аннотирования; ③ после аннотирования 500 образцов разрыв между тремя методами сокращается (93,0–93,81 TP3T), а при полном объеме аннотирования в 1 500 образцов результаты практически сравняются (94,0–94,31 TP3T); ④ MoCo v2 незначительно превосходит SimCLR при любом объеме аннотаций (+0,6–1,81 TP3T) благодаря стабильности кодировщика импульса.

Подробное сравнение SimCLR и MoCo v2

Рекомендации по промышленному внедрению: если объем видеопамяти GPU достаточен (≥24 ГБ, можно установить Batch ≥256) и требуется быстрая итерация, выбирайте SimCLR (быстрое обучение, простой код). Если объем видеопамяти GPU ограничен (≤16 ГБ) или требуется максимальная точность, выбирайте MoCo v2. При внедрении в сценарии с крановыми системами рекомендуется использовать MoCo v2 — поскольку максимальный размер партии на одной карте Jetson AGX Orin (64 ГБ) составляет примерно 128, а MoCo v2 лучше подходит для работы с небольшими партиями.

Эксперимент по сравнению стратегий расширения данных

Расширение данных является ключом к успеху контрастного обучения. В данном эксперименте, используя настройки MoCo v2+100, мы поочередно отключали каждый из компонентов расширения данных, чтобы проверить их вклад:

Случайная обрезка является ключевым улучшением в методе контрастного обучения — она заставляет модель изучать локальные текстурные паттерны объекта, а не полагаться на общую композицию изображения (в промышленных сценах дефекты обычно занимают лишь 5–20% изображения). На втором месте находится цветовая дрожание, поскольку цветовая гамма изображений промышленных деталей ограничена (преимущественно оттенки серого или металлические цвета), но дрожание повышает устойчивость к изменениям освещения. Гауссово размытие служит для фильтрации высокочастотного шума и улучшает способность модели к обобщению текстур поверхности.

Реализация проекта — исходный код ядра MoCo v2

Обновление кодировщика импульса # MoCo v2 (псевдокод на PyTorch)
class MoCo(nn.Module):
    def __init__(self, base_encoder=resnet18, dim=128, K=4096, m=0.999, T=0.2):
 super().__init__()
        self.K = K  # длина очереди
 self.m = m  # коэффициент импульса
 self.T = T  # коэффициент температуры
 self.encoder_q = base_encoder()  # кодировщик запросов (обучение)
        self.encoder_k = base_encoder()  # кодировщик ключей (обновление импульса)
 for p in self.encoder_k.parameters():
 p.requires_grad = False  # заморозить кодировщик импульса
        # Очередь: хранение K характеристик отрицательных образцов
 self.register_buffer("queue", torch.randn(dim, K))
 self.queue = F.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

 @torch.no_grad()
    def _momentum_update(self):
        # Обновление импульса: θ_k = m*θ_k + (1-m)*θ_q
 for p_q, p_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            p_k.data = self.m * p_k.data + (1 - self.m) * p_q.data

 def forward(self, im_q, im_k):
 q = self.encoder_q(im_q)  # Запрос характеристик NxC
        q = F.normalize(q, dim=1)
 with torch.no_grad():
 self._momentum_update()
 k = self.encoder_k(im_k)   # ключевые характеристики NxC
            k = F.normalize(k, dim=1)
 # Потеря сравнения: l_pos — пары положительных образцов + l_neg — пары отрицательных образцов
 l_pos = torch.einsum("nc,nc->n", q, k).unsqueeze(-1)  # Nx1
 l_neg = torch.einsum("nc,ck->nk", q, self.queue.clone().detach())  # NxK
        logits = torch.cat([l_pos, l_neg], dim=1) / self.T  # Nx(1+K)
 labels = torch.zeros(logits.shape[0], dtype=torch.long)
        return F.cross_entropy(logits, labels)

Часто задаваемые вопросы