#!/usr/bin/env python3 # -*- coding:utf-8 -*- # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. import torch import torch.distributed as dist import torch.nn as nn import os import random from .base_exp import BaseExp class Exp(BaseExp): def __init__(self): super().__init__() # ---------------- model config ---------------- # self.num_classes = 80 self.depth = 1.00 self.width = 1.00 # ---------------- dataloader config ---------------- # # set worker to 4 for shorter dataloader init time self.data_num_workers = 4 self.input_size = (640, 640) self.random_size = (14, 26) self.train_ann = "instances_train2017.json" self.val_ann = "instances_val2017.json" # --------------- transform config ----------------- # self.degrees = 10.0 self.translate = 0.1 self.scale = (0.1, 2) self.mscale = (0.8, 1.6) self.shear = 2.0 self.perspective = 0.0 self.enable_mixup = True # -------------- training config --------------------- # self.warmup_epochs = 5 self.max_epoch = 300 self.warmup_lr = 0 self.basic_lr_per_img = 0.01 / 64.0 self.scheduler = "yoloxwarmcos" self.no_aug_epochs = 15 self.min_lr_ratio = 0.05 self.ema = True self.weight_decay = 5e-4 self.momentum = 0.9 self.print_interval = 10 self.eval_interval = 10 self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] # ----------------- testing config ------------------ # self.test_size = (640, 640) self.test_conf = 0.001 self.nmsthre = 0.65 def get_model(self): from yolox.models import YOLOPAFPN, YOLOX, YOLOXHead def init_yolo(M): for m in M.modules(): if isinstance(m, nn.BatchNorm2d): m.eps = 1e-3 m.momentum = 0.03 if getattr(self, "model", None) is None: in_channels = [256, 512, 1024] backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels) head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels) self.model = YOLOX(backbone, head) self.model.apply(init_yolo) self.model.head.initialize_biases(1e-2) return self.model def get_data_loader(self, batch_size, is_distributed, no_aug=False): from yolox.data import ( COCODataset, DataLoader, InfiniteSampler, MosaicDetection, TrainTransform, YoloBatchSampler ) dataset = COCODataset( data_dir=None, json_file=self.train_ann, img_size=self.input_size, preproc=TrainTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=50, ), ) dataset = MosaicDetection( dataset, mosaic=not no_aug, img_size=self.input_size, preproc=TrainTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=120, ), degrees=self.degrees, translate=self.translate, scale=self.scale, shear=self.shear, perspective=self.perspective, enable_mixup=self.enable_mixup, ) self.dataset = dataset if is_distributed: batch_size = batch_size // dist.get_world_size() sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0) batch_sampler = YoloBatchSampler( sampler=sampler, batch_size=batch_size, drop_last=False, input_dimension=self.input_size, mosaic=not no_aug, ) dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True} dataloader_kwargs["batch_sampler"] = batch_sampler train_loader = DataLoader(self.dataset, **dataloader_kwargs) return train_loader def random_resize(self, data_loader, epoch, rank, is_distributed): tensor = torch.LongTensor(2).cuda() if rank == 0: size_factor = self.input_size[1] * 1.0 / self.input_size[0] size = random.randint(*self.random_size) size = (int(32 * size), 32 * int(size * size_factor)) tensor[0] = size[0] tensor[1] = size[1] if is_distributed: dist.barrier() dist.broadcast(tensor, 0) input_size = data_loader.change_input_dim( multiple=(tensor[0].item(), tensor[1].item()), random_range=None ) return input_size def get_optimizer(self, batch_size): if "optimizer" not in self.__dict__: if self.warmup_epochs > 0: lr = self.warmup_lr else: lr = self.basic_lr_per_img * batch_size pg0, pg1, pg2 = [], [], [] # optimizer parameter groups for k, v in self.model.named_modules(): if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter): pg2.append(v.bias) # biases if isinstance(v, nn.BatchNorm2d) or "bn" in k: pg0.append(v.weight) # no decay elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter): pg1.append(v.weight) # apply decay optimizer = torch.optim.SGD( pg0, lr=lr, momentum=self.momentum, nesterov=True ) optimizer.add_param_group( {"params": pg1, "weight_decay": self.weight_decay} ) # add pg1 with weight_decay optimizer.add_param_group({"params": pg2}) self.optimizer = optimizer return self.optimizer def get_lr_scheduler(self, lr, iters_per_epoch): from yolox.utils import LRScheduler scheduler = LRScheduler( self.scheduler, lr, iters_per_epoch, self.max_epoch, warmup_epochs=self.warmup_epochs, warmup_lr_start=self.warmup_lr, no_aug_epochs=self.no_aug_epochs, min_lr_ratio=self.min_lr_ratio, ) return scheduler def get_eval_loader(self, batch_size, is_distributed, testdev=False): from yolox.data import COCODataset, ValTransform valdataset = COCODataset( data_dir=None, json_file=self.val_ann if not testdev else "image_info_test-dev2017.json", name="val2017" if not testdev else "test2017", img_size=self.test_size, preproc=ValTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225) ), ) if is_distributed: batch_size = batch_size // dist.get_world_size() sampler = torch.utils.data.distributed.DistributedSampler( valdataset, shuffle=False ) else: sampler = torch.utils.data.SequentialSampler(valdataset) dataloader_kwargs = { "num_workers": self.data_num_workers, "pin_memory": True, "sampler": sampler, } dataloader_kwargs["batch_size"] = batch_size val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs) return val_loader def get_evaluator(self, batch_size, is_distributed, testdev=False): from yolox.evaluators import COCOEvaluator val_loader = self.get_eval_loader(batch_size, is_distributed, testdev=testdev) evaluator = COCOEvaluator( dataloader=val_loader, img_size=self.test_size, confthre=self.test_conf, nmsthre=self.nmsthre, num_classes=self.num_classes, testdev=testdev, ) return evaluator def eval(self, model, evaluator, is_distributed, half=False): return evaluator.evaluate(model, is_distributed, half)