BiFPN

Thêm lớp bifpn vào common.py

BiFPN (Bi-directional Feature Pyramid Network) là một kiến trúc mạng kim tự tháp đặc trưng được sử dụng cho các nhiệm vụ phát hiện mục tiêu.

Trong các thuật toán phát hiện mục tiêu truyền thống, thường sử dụng FPN (Feature Pyramid Network) để trích xuất các đặc trưng ở nhiều tỷ lệ khác nhau. FPN đạt được điều này bằng cách thêm các nhánh phụ và hoạt động tăng kích thước vào mạng. Tuy nhiên, trong FPN, dòng thông tin giữa các hoạt động tăng và giảm kích thước chỉ là đơn hướng.

BiFPN khắc phục vấn đề này bằng cách giới thiệu kết nối hai chiều giữa các tầng của mạng kim tự tháp đặc trưng ở các tỷ lệ khác nhau, cho phép trao đổi và hợp nhất thông tin giữa các tầng. Cụ thể, BiFPN bao gồm các bước sau:

1. Tăng kích thước kim tự tháp: Bắt đầu từ các bản đồ đặc trưng ở mức độ thấp hơn và tạo ra các bản đồ đặc trưng có độ phân giải cao hơn thông qua hoạt động tăng kích thước.
2. Giảm kích thước kim tự tháp: Bắt đầu từ các bản đồ đặc trưng ở mức độ cao hơn và tạo ra các bản đồ đặc trưng có độ phân giải thấp hơn nhưng giàu thông tin ngữ nghĩa hơn thông qua hoạt động giảm kích thước.
3. Kết nối hai chiều: Kết nối các bản đồ đặc trưng từ quá trình tăng và giảm kích thước để thực hiện dòng thông tin hai chiều.
4. Hợp nhất đặc trưng: Hợp nhất các bản đồ đặc trưng từ các tầng liền kề để các đặc trưng ở các tỷ lệ khác nhau có thể bổ sung cho nhau, cải thiện độ chính xác của việc phát hiện mục tiêu.

class FusionNetwork(nn.Module):
    def __init__(self, input_channels, output_channels):
        super(FusionNetwork, self).__init__()
        self.weights = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
        self.smoothing = 0.0001
        self.transformation = nn.Conv2d(input_channels, output_channels, kernel_size=1, stride=1, padding=0)
        self.activation = nn.SiLU()

    def forward(self, inputs):
        weighted_inputs = []
        for i in range(len(inputs)):
            weighted_inputs.append(inputs[i] * (self.weights[i] / (torch.sum(self.weights) + self.smoothing)))
        fused_output = torch.sum(torch.stack(weighted_inputs), dim=0)
        transformed_output = self.transformation(fused_output)
        activated_output = self.activation(transformed_output)
        return activated_output

Lớp FusionNetwork kế thừa từ lớp nn.Module của PyTorch. Trong hàm khởi tạo __init__, nó định nghĩa các trọng số học được (self.weights) và một hằng số nhỏ (self.smoothing) để đảm bảo ổn định số học. Sau đó, nó áp dụng một lớp tích chập (self.transformation) và hàm kích hoạt SiLU (self.activation) lên dữ liệu đầu vào.

Nhắc lại về nn.Module

Trong PyTorch, nn.Module là lớp cơ sở dùng để định nghĩa mô hình mạng nơ-ron. Bằng cách kế thừa từ nn.Module, chúng ta có thể tạo ra các mô hình tùy chỉnh, bao gồm cả việc định nghĩa cấu trúc mạng và triển khai quá trình truyền xuôi.

Các chức năng chính của nn.Module bao gồm:

1. Khởi tạo mô hình: Định nghĩa các tầng và tham số học được trong hàm khởi tạo __init__.
2. Triển khai quá trình truyền xuôi: Định nghĩa hàm forward để mô tả cách thức tính toán đầu ra từ đầu vào.
3. Quản lý tham số mô hình: Cung cấp các phương pháp như parameters() để lấy danh sách tất cả các tham số học được và to() để di chuyển mô hình đến thiết bị cụ thể (như GPU).

Thêm cấu trúc fusion vào yolo.py

# Thêm cấu trúc fusion
elif module_type in [FusionNetwork]:
    output_channels = max([channels[x] for x in from_layers])

Mã nguồn trên nằm trong hàm parse_model của YOLO, chịu trách nhiệm phân tích từ điển định nghĩa mô hình và xây dựng các tầng cũng như tham số tương ứng.

Chỉnh sửa train.py

standard_batch_size = 64
accumulate_loss = max(round(standard_batch_size / current_batch_size), 1)
hyperparameters['weight_decay'] *= current_batch_size * accumulate_loss / standard_batch_size
LOGGER.info(f"Weight decay đã được scale: {hyperparameters['weight_decay']}")

group_1, group_2, group_3 = [], [], []  # Nhóm tham số tối ưu hóa
for layer in model.modules():
    if hasattr(layer, 'bias') and isinstance(layer.bias, nn.Parameter):  # bias
        group_3.append(layer.bias)
    if isinstance(layer, nn.BatchNorm2d):  # weight (không decay)
        group_1.append(layer.weight)
    elif hasattr(layer, 'weight') and isinstance(layer.weight, nn.Parameter):  # weight (có decay)
        group_2.append(layer.weight)
    elif isinstance(layer, FusionNetwork) and hasattr(layer, 'weights') and isinstance(layer.weights, nn.Parameter):
        group_2.append(layer.weights)

Mã nguồn trên chỉnh sửa cách xử lý các tham số của mô hình để tối ưu hóa hiệu quả hơn.

Thay thế Concat bằng Fusion

nc: 80  # Số lượng lớp
depth_multiple: 0.33  # Hệ số nhân độ sâu mô hình
width_multiple: 0.50  # Hệ số nhân kênh tầng
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# Backbone YOLOv5 v6.0
backbone:
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9

# Đầu YOLOv5 v6.1 với BiFPN
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, FusionNetwork, [256, 256]],  # Ghép backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, FusionNetwork, [128, 128]],  # Ghép backbone P3
   [-1, 3, C3, [256, False]],  # 17 

   [-1, 1, Conv, [512, 3, 2]],  
   [[-1, 13, 6], 1, FusionNetwork, [256, 256]],  # Phiên bản v5s có số kênh là nửa giá trị mặc định
   [-1, 3, C3, [512, False]],  # 20 

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, FusionNetwork, [256, 256]],  # Ghép head P5
   [-1, 3, C3, [1024, False]],  # 23 

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Phát hiện (P3, P4, P5)
  ]