關於yolov8訓練的一些改動及注意事項

2023-02-06 06:01:31

1、YOLOv8創新改進點：

1.1.Backbone

使用的依舊是CSP的思想，不過YOLOv5中的C3模組被替換成了C2f模組，實現了進一步的輕量化，同時YOLOv8依舊使用了YOLOv5等架構中使用的SPPF模組；

1.2.PAN-FPN

毫無疑問YOLOv8依舊使用了PAN的思想，不過通過對比YOLOv5與YOLOv8的結構圖可以看到，YOLOv8將YOLOv5中PAN-FPN上取樣階段中的折積結構刪除了，同時也將C3模組替換為了C2f模組

1.3.Decoupled-Head

是不是嗅到了不一樣的味道？是的，YOLOv8走向了Decoupled-Head；

1.4.Anchor-Free

YOLOv8拋棄了以往的Anchor-Base，使用了Anchor-Free的思想；

1.5.損失函數

YOLOv8使用VFL Loss作為分類損失，使用DFL Loss+CIOU Loss作為分類損失；

1.6.樣本匹配

YOLOv8拋棄了以往的IOU匹配或者單邊比例的分配方式，而是使用了Task-Aligned Assigner匹配方式。

2、關於基於預訓練模型的訓練

yolov8版本更新後，程式碼結構也隨著更新，跟v5的結構大不一樣，大部分介面以及網路結構也隨之改動，為了加速演演算法落地，我們在訓練時一般會遷移一部分預訓練引數從而是的模型達到較好的效果，但是若你的模型跟預訓練模型只有一小部分相似，但是又想繼承這一小部分的特徵，直接載入所有引數訓練肯定是不可取的，那就需要進行神經網路的層凍結，通過凍結一些層來使得模型加速擬合，減少引數訓練量。例如：當你的網路很複雜，他的前端網路是一個 vgg-16 的分類網路，後面要拼接一個自己寫的功能網路，這個時候，你把 vgg-16 的網路架構定義好了之後，上網下載vgg-16 的訓練好的網路引數，然後載入到你寫的網路中，然後把 vgg-16 相關的層凍結掉，只訓練你自己寫的小網路的引數。這樣的話，你就可以省掉很多的運算資源和時間，提高效率。

注意：凍結網路層之後，最好對網路重新 compile 一下，否則在一些場景下不會生效，compile 才會生效。

廢話不多說了，上乾貨：

def _setup_train(self, rank, world_size):
        """
        Builds dataloaders and optimizer on correct rank process.
        """
        # model
        self.run_callbacks("on_pretrain_routine_start")
        ckpt = self.setup_model()
        self.model = self.model.to(self.device)
        freeze=[5]
        freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
        for k, v in self.model.named_parameters():
            v.requires_grad = True  # train all layers
            # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
            if any(x in k for x in freeze):
                LOGGER.info(f'freezing {k}')
                v.requires_grad = False
        self.set_model_attributes()
        if world_size > 1:
            self.model = DDP(self.model, device_ids=[rank])
        # Check imgsz
        gs = max(int(self.model.stride.max() if hasattr(self.model, 'stride') else 32), 32)  # grid size (max stride)
        self.args.imgsz = check_imgsz(self.args.imgsz, stride=gs, floor=gs)
        # Batch size
        if self.batch_size == -1:
            if RANK == -1:  # single-GPU only, estimate best batch size
                self.batch_size = check_train_batch_size(self.model, self.args.imgsz, self.amp)
            else:
                SyntaxError('batch=-1 to use AutoBatch is only available in Single-GPU training. '
                            'Please pass a valid batch size value for Multi-GPU DDP training, i.e. batch=16')

        # Optimizer
        self.accumulate = max(round(self.args.nbs / self.batch_size), 1)  # accumulate loss before optimizing
        self.args.weight_decay *= self.batch_size * self.accumulate / self.args.nbs  # scale weight_decay
        self.optimizer = self.build_optimizer(model=self.model,
                                              name=self.args.optimizer,
                                              lr=self.args.lr0,
                                              momentum=self.args.momentum,
                                              decay=self.args.weight_decay)
        # Scheduler
        if self.args.cos_lr:
            self.lf = one_cycle(1, self.args.lrf, self.epochs)  # cosine 1->hyp['lrf']
        else:
            self.lf = lambda x: (1 - x / self.epochs) * (1.0 - self.args.lrf) + self.args.lrf  # linear
        self.scheduler = lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
        self.scheduler.last_epoch = self.start_epoch - 1  # do not move
        self.stopper, self.stop = EarlyStopping(patience=self.args.patience), False

        # dataloaders
        batch_size = self.batch_size // world_size if world_size > 1 else self.batch_size
        self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=rank, mode="train")
        if rank in {0, -1}:
            self.test_loader = self.get_dataloader(self.testset, batch_size=batch_size * 2, rank=-1, mode="val")
            self.validator = self.get_validator()
            metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix="val")
            self.metrics = dict(zip(metric_keys, [0] * len(metric_keys)))  # TODO: init metrics for plot_results()?
            self.ema = ModelEMA(self.model)
        self.resume_training(ckpt)
        self.run_callbacks("on_pretrain_routine_end")