DETR

Изначально имеем набор боксов на изображении. Далее:

1. Сортируем bbox-ы, опираясь на их confidence scores.
2. Выбираем bbox c максимальным confidence score, сохраняем его как итоговое предсказание детектора и удаляем из первоначального набора. Таких боксов будет один или несколько (об этом расскажем чуть позже).
3. Удаляем из набора все боксы, которые имеют пересечение с боксом / боксами итогового предсказания выше заданного порога. Пересечение считается, например, как Intersection over Union.
4. Повторяем шаги 2 и 3, пока в исходном наборе не останется боксов.

from typing import List


def non_max_suppression(
		boxes: List[List[int]],
		scores: List[float],
		threshold: float
):
    """
    Алгоритм NMS для отбора детекций из набора боксов с соответсвующими им
    confidence scores.
   
    :param boxes: набор bbox-ов формата [xmin, ymin, xmax, ymax]
    :param scores: список confidence scores
    :param threshold: IoU (intersection-over-union) threshold
    :return: список индексов боксов из исходного набора, которые останутся
    после NMS.
    """
    # 1. Сортируем bbox-ы, опираясь на их confidence scores
    order = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    keep = []
    while order:
		    # 2. Выбираем bbox c максимальным confidence score, сохраняем его как
		    # итоговое предсказание детектора и удаляем из первоначального набора.
        i = order.pop(0)
        keep.append(i)
        for j in order:
            # 3. Считаем IoU между итоговым предсказанием и остальными боксами.
            intersection = max(0, min(boxes[i][2], boxes[j][2]) - max(boxes[i][0], boxes[j][0])) * \
								           max(0, min(boxes[i][3], boxes[j][3]) - max(boxes[i][1], boxes[j][1]))
            union = (boxes[i][2] - boxes[i][0]) * (boxes[i][3] - boxes[i][1]) + \
                    (boxes[j][2] - boxes[j][0]) * (boxes[j][3] - boxes[j][1]) - intersection
            iou = intersection / union

            # Удаляем все боксы из набора, которые имеют пересечение
            # выше заданного порога.
            if iou > threshold:
                order.remove(j)
    return keep

import torch
import torch.nn as nn


class TransformerEncoderLayer(nn.Module):
		"""
		Имплементация слоя энкодера трансформера.
    """
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Слои для FFN
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
				
				# Функция для выбора активации из списка: [ReLU, GELU, GLU]
        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,
                     src,
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        # Forward для случая, когда нормализация идёт после MHSA и MLP.
        # Добавляем позиционное кодирование только к Query и Key.
        q = k = self.with_pos_embed(src, pos)
        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

    def forward_pre(self, src,
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        # Forward для случая, когда нормализация идёт до MHSA и MLP.
        src2 = self.norm1(src)
        # Добавляем позиционное кодирование только к Query и Key.
        q = k = self.with_pos_embed(src2, pos)
        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)
        return src

    def forward(self, src,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)

import torch
import torch.nn as nn


class TransformerEncoder(nn.Module):
		"""
		Имплементация энкодера трансформера.
    """
    def __init__(self, encoder_layer, num_layers, norm=None):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        output = src
				
				# Стакаем N раз энкодер слои.
        for layer in self.layers:
            output = layer(output, src_mask=mask,
                           src_key_padding_mask=src_key_padding_mask, pos=pos)

        if self.norm is not None:
            output = self.norm(output)

        return output

import torch
import torch.nn as nn


class TransformerDecoderLayer(nn.Module):
		"""
		Имплементация слоя декодера трансформера.
    """
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Слои для FFN
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
				
				# Функция для выбора активации из списка: [ReLU, GELU, GLU]
        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,
                     tgt_mask: Optional[Tensor] = None,
                     memory_mask: Optional[Tensor] = None,
                     tgt_key_padding_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):
        # Forward для случая, когда нормализация идёт после MHSA и MLP.
        # Добавляем позиционное кодирование только к Query и Key.
        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        # Обратите внимание на то, что подаётся как Q, K, V в attention.
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt

    def forward_pre(self, tgt, memory,
                    tgt_mask: Optional[Tensor] = None,
                    memory_mask: Optional[Tensor] = None,
                    tgt_key_padding_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        # Forward для случая, когда нормализация идёт до MHSA и MLP.
        tgt2 = self.norm1(tgt)
        # Добавляем позиционное кодирование только к Query и Key.
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)
        # Обратите внимание на то, что подаётся как Q, K, V в attention.
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)

import torch
import torch.nn as nn


class TransformerDecoder(nn.Module):
		"""
		Имплементация декодера трансформера.
    """
    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
        super().__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.return_intermediate = return_intermediate

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        output = tgt

        intermediate = []
				
				# Стакаем M раз декодер слои и записываем их в список intermediate.
        for layer in self.layers:
            output = layer(output, memory, tgt_mask=tgt_mask,
                           memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask,
                           memory_key_padding_mask=memory_key_padding_mask,
                           pos=pos, query_pos=query_pos)
            if self.return_intermediate:
                intermediate.append(self.norm(output))

        if self.norm is not None:
            output = self.norm(output)
            if self.return_intermediate:
                intermediate.pop()
                intermediate.append(output)

        if self.return_intermediate:
            return torch.stack(intermediate)

        return output.unsqueeze(0)

import torch
import torch.nn as nn


class Transformer(nn.Module):
		"""
		Имплементация трансформера.
    """
    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=False):
        super().__init__()
				
				# Инициализируем энкодер.
        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
				
				# Инициализируем декодер.
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=return_intermediate_dec)

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, src, mask, query_embed, pos_embed):
        # Разворачиваем входной тензор размера NxCxHxW в тензор размера HWxNxC
        bs, c, h, w = src.shape
        src = src.flatten(2).permute(2, 0, 1)
        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
        mask = mask.flatten(1)

        tgt = torch.zeros_like(query_embed)
        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
        hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                          pos=pos_embed, query_pos=query_embed)
        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)

import torch
import torch.nn as nn


class MLP(nn.Module):
    """
		Имплементация MLP.
    """
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        # output dim = 4, так как bbox = [x, y, w, h]
        # num_layers = 3
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        # если input_dim = hidden_dim = 512, то размерности 3-х Linear слоёв - 
	      # (512, 512), (512, 512), (512, 4)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

import torch
import torch.nn as nn


def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
	  """
	  Вспомогательная функция.
	  """
    if tensor_list[0].ndim == 3:
        if torchvision._is_tracing():
            # оптимизация для ONNX
            return _onnx_nested_tensor_from_tensor_list(tensor_list)

        # support different-sized images
        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
        batch_shape = [len(tensor_list)] + max_size
        b, c, h, w = batch_shape
        dtype = tensor_list[0].dtype
        device = tensor_list[0].device
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
            m[: img.shape[1], :img.shape[2]] = False
    else:
        raise ValueError('not supported')
    return NestedTensor(tensor, mask)


class DETR(nn.Module):
    """
		Имплементация DETR.
    """
    def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):
        """
        Инициализация модели.
        Parameters:
            backbone: CNN бэкбон.
            transformer: Encoder-Decoder transformer.
            num_queries: количество object queries. это максимальное число
	            детекций DETR на одно изображение. Для COCO авторы
	            рекомендуют брать число 100.
            
            aux_loss: True если используется вспомогательный лосс для декодинга.
        """
        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        # num_classes + 1, потому что мы не забываем про no_objects
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        self.query_embed = nn.Embedding(num_queries, hidden_dim)
        self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)
        self.backbone = backbone
        self.aux_loss = aux_loss

    def forward(self, samples: NestedTensor):
        """
        На вход forward принимает вложенный тензор, который состоит из:
        The forward expects a NestedTensor, which consists of:
             - samples.tensor: батч картинок размера [batch_size x 3 x H x W].
             - samples.mask: бинарная маска размера [batch_size x H x W],
             содержащая 1 на padded пикселях.
						
			Forward возвращает dict со следующими элементами:
             - "pred_logits": классификационные логиты, включая no-object,
             для всех queries.
              Размер = [batch_size x num_queries x (num_classes + 1)]
              
             - "pred_boxes": нормализованные координаты боксов 
             значения от 0 до 1) для всех queries размера
              (center_x, center_y, height, width).
              
             - "aux_outputs": если aux_loss == True, возвращает их значения.
        """
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)

        src, mask = features[-1].decompose()
        assert mask is not None
        hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]
				
				# классификационная голова
        outputs_class = self.class_embed(hs)
        # регрессионная голова
        outputs_coord = self.bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        return out

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

Предположим, есть следующая задача:

У нас есть 3 сотрудника, которые должны доехать до 3-х разных клиентов. Мы знаем стоимость поездки каждого сотрудника до каждого клиента в вечно зеленых [долларах]. Как оптимально (с наименьшими затратами) распределить сотрудников по клиентам или, что то же самое, распределить поездки (jobs) по сотрудникам?

В самом начале нам нужно посчитать матрицу стоимости. Именно с ней и будет работать алгоритм.

После того, как матрица найдена, Венгерский алгоритм предлагает нахождение решения за 5 шагов:

• Вычитаем самое маленькое значение в строке из всех ее элементов для каждой строки (самое маленькое значение в каждой строке при этом станет равно 0).

• Вычитаем самое маленькое значение в столбце из всех его элементов для каждого столбца (самое маленькое значение в каждом столбце при этом станет равно 0).

• Зачеркиваем все нули минимальным количеством линий. Если количество линий получилось меньше n, где n — количество сотрудников / работ, то переходим к шагу 4. В противном случае — переходим к шагу 5.

• Находим наименьший элемент, не охваченный ни одной линией, и вычитаем его из всей матрицы. Если элемент был охвачен какой-либо линией дважды — добавляем в то место, где он дважды перечеркнут. Затем возвращаемся к шагу 3.

• Назначаем работы сотрудникам, начиная со строки только с одним нулем. Каждый раз, когда мы сопоставляем одну работу с сотрудником, мы пересекаем его строку и столбец, чтобы сделать его недоступным.

Таким образом, в нашей задаче минимальное количество долларов, которое нам потребуется — это (всего лишь 😉) 26.

import torch
import torch.nn as nn


class MHAttentionMap(nn.Module):
    """
		Имплементация MHAttentionMap.
    """
    def __init__(self, query_dim, hidden_dim, num_heads, dropout=0.0, bias=True):
        super().__init__()
        self.num_heads = num_heads
        self.hidden_dim = hidden_dim
        self.dropout = nn.Dropout(dropout)

        self.q_linear = nn.Linear(query_dim, hidden_dim, bias=bias)
        self.k_linear = nn.Linear(query_dim, hidden_dim, bias=bias)
				
				# Инициализация весов.
        nn.init.zeros_(self.k_linear.bias)
        nn.init.zeros_(self.q_linear.bias)
        nn.init.xavier_uniform_(self.k_linear.weight)
        nn.init.xavier_uniform_(self.q_linear.weight)
        self.normalize_fact = float(hidden_dim / self.num_heads) ** -0.5

    def forward(self, q, k, mask: Optional[Tensor] = None):
        q = self.q_linear(q)
        k = F.conv2d(k, self.k_linear.weight.unsqueeze(-1).unsqueeze(-1), self.k_linear.bias)
        qh = q.view(q.shape[0], q.shape[1], self.num_heads, self.hidden_dim // self.num_heads)
        kh = k.view(k.shape[0], self.num_heads, self.hidden_dim // self.num_heads, k.shape[-2], k.shape[-1])
        weights = torch.einsum("bqnc,bnchw->bqnhw", qh * self.normalize_fact, kh)

        if mask is not None:
            weights.masked_fill_(mask.unsqueeze(1).unsqueeze(1), float("-inf"))
        weights = F.softmax(weights.flatten(2), dim=-1).view(weights.size())
        weights = self.dropout(weights)
        return weights

import torch
import torch.nn as nn


class MaskHeadSmallConv(nn.Module):
		"""
		Имплементация MaskHeadSmallConv. Upsampling делается с помощью FPN подхода.
    """
    def __init__(self, dim, fpn_dims, context_dim):
        super().__init__()
				
				# В качестве нормализации активно используем GroupNorm.
				# Идея в том, чтобы построить feature map в разных масштабах: 
				# от большего к меньшему, которые затем мы будем складывать друг
				# с другом.
        inter_dims = [dim, context_dim // 2, context_dim // 4, context_dim // 8, context_dim // 16, context_dim // 64]
        self.lay1 = torch.nn.Conv2d(dim, dim, 3, padding=1)
        self.gn1 = torch.nn.GroupNorm(8, dim)
        self.lay2 = torch.nn.Conv2d(dim, inter_dims[1], 3, padding=1)
        self.gn2 = torch.nn.GroupNorm(8, inter_dims[1])
        self.lay3 = torch.nn.Conv2d(inter_dims[1], inter_dims[2], 3, padding=1)
        self.gn3 = torch.nn.GroupNorm(8, inter_dims[2])
        self.lay4 = torch.nn.Conv2d(inter_dims[2], inter_dims[3], 3, padding=1)
        self.gn4 = torch.nn.GroupNorm(8, inter_dims[3])
        self.lay5 = torch.nn.Conv2d(inter_dims[3], inter_dims[4], 3, padding=1)
        self.gn5 = torch.nn.GroupNorm(8, inter_dims[4])
        self.out_lay = torch.nn.Conv2d(inter_dims[4], 1, 3, padding=1)

        self.dim = dim

        self.adapter1 = torch.nn.Conv2d(fpn_dims[0], inter_dims[1], 1)
        self.adapter2 = torch.nn.Conv2d(fpn_dims[1], inter_dims[2], 1)
        self.adapter3 = torch.nn.Conv2d(fpn_dims[2], inter_dims[3], 1)
				
				# Инициализация весов.
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_uniform_(m.weight, a=1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x: Tensor, bbox_mask: Tensor, fpns: List[Tensor]):
		    """
		    fpns - 3 feature map разного размера с бэкбона, которые ниже
		    мы будем склеивать с feature maps, полученными свёртками в mask head.
		    """
		    
		    # Мы объединяем x (feature map с backbone) и масками боксов,
		    # полученных на этапе MHAttentionMap. Размеры итогового тензора
		    # будут зависеть от количества боксов.
        x = torch.cat([_expand(x, bbox_mask.shape[1]), bbox_mask.flatten(0, 1)], 1)

        x = self.lay1(x)
        x = self.gn1(x)
        x = F.relu(x)
        x = self.lay2(x)
        x = self.gn2(x)
        x = F.relu(x)
				
				# Строим пирамиду из признаков.
        cur_fpn = self.adapter1(fpns[0])
        if cur_fpn.size(0) != x.size(0):
		        # expand - resize feature map cur_fpn до размеров feature map x.
            cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
        x = cur_fpn + F.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
        x = self.lay3(x)
        x = self.gn3(x)
        x = F.relu(x)

        cur_fpn = self.adapter2(fpns[1])
        if cur_fpn.size(0) != x.size(0):
            cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
        x = cur_fpn + F.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
        x = self.lay4(x)
        x = self.gn4(x)
        x = F.relu(x)

        cur_fpn = self.adapter3(fpns[2])
        if cur_fpn.size(0) != x.size(0):
            cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
        x = cur_fpn + F.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
        x = self.lay5(x)
        x = self.gn5(x)
        x = F.relu(x)

        x = self.out_lay(x)
        return x

import torch
import torch.nn as nn


class DETRsegm(nn.Module):
		"""
		Имплементация сегментационного DETR.
    """
    def __init__(self, detr, freeze_detr=False):
        super().__init__()
        self.detr = detr
				
				# Как было сказано выше, мы можем просто заморозить веса детектора.
				# Это не повлияет на итоговое качество, но существенно упростит 
				# обучение.
        if freeze_detr:
            for p in self.parameters():
                p.requires_grad_(False)

        hidden_dim, nheads = detr.transformer.d_model, detr.transformer.nhead
        self.bbox_attention = MHAttentionMap(hidden_dim, hidden_dim, nheads, dropout=0.0)
        self.mask_head = MaskHeadSmallConv(hidden_dim + nheads, [1024, 512, 256], hidden_dim)

    def forward(self, samples: NestedTensor):
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.detr.backbone(samples)

        bs = features[-1].tensors.shape[0]

        src, mask = features[-1].decompose()
        assert mask is not None
        src_proj = self.detr.input_proj(src)
        hs, memory = self.detr.transformer(src_proj, mask, self.detr.query_embed.weight, pos[-1])

        outputs_class = self.detr.class_embed(hs)
        outputs_coord = self.detr.bbox_embed(hs).sigmoid()
        out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord[-1]}
        if self.detr.aux_loss:
            out['aux_outputs'] = self.detr._set_aux_loss(outputs_class, outputs_coord)

        # MHSA формирует attention maps для боксов.
        bbox_mask = self.bbox_attention(hs[-1], memory, mask=mask)
				
				# А mask head, которая и учит маски объектов.
        seg_masks = self.mask_head(src_proj, bbox_mask, [features[2].tensors, features[1].tensors, features[0].tensors])
        outputs_seg_masks = seg_masks.view(bs, self.detr.num_queries, seg_masks.shape[-2], seg_masks.shape[-1])

        out["pred_masks"] = outputs_seg_masks
        return out


def _expand(tensor, length: int):
    return tensor.unsqueeze(1).repeat(1, int(length), 1, 1, 1).flatten(0, 1)