[Detectron2] code level review

2021. 9. 28. 14:07AI 부스트캠프

Detectron2 ?

def. Facebook에서 만든 pytorch 기반으로 하고 Object detection/sementic segmentation을 위한 training/inference opensorce libary

 

 

Why Detectron2 ?

python에 최적화가 되어있다!

연산량이 많이 드는 코드는 C를 통해서 구현했다

box IOU, defromable convolution 부분등은 CUDA로 구현했다

 

 

Detectron2 Structure

  • tools dir
    • plain_train_net.py : 구조파악이 쉽지만, SGD를 이용한 학습만 지원하고 나머지 기능들은 지원하지 않는다
    • train_net.py : training iter가 추상화가 잘 되있기 때문에 실제 디버깅하는 부분에서 힘들다. 하지만 Engine을 이용한 학습을 하며, 이 Engine을 이용하면 사용자는 학습과정을 신경쓰지않고 모델과 loss에만 집중할 수 있다.

 

 

  • detectron2 dir
    • config folder : 필요한 하이퍼 파라미터를 정의함 (수정가능)
    • engine folder : caffe 스타일의 training code를 pytorch로 구현함, 이를 통해 학습을 추상화 시킴
    • model_zoo folder : Facebook에서 학습한 모델들이 관리되는 폴더
    • layers folder : python으로 구현하면 time complexitiy가 높아 C와 CUDA로 구현해야되는 부분들

 

 

 

Code level Analysis

글쓴이도 Decetron2를 써보는 것이 처음이므로 최대한 한줄한줄 풀어서 설명해보려고 한다

 

try:
    register_coco_instances('coco_trash_train', {}, '../dataset/train.json', '../dataset/')
except AssertionError:
    pass

try:
    register_coco_instances('coco_trash_test', {}, '../dataset/test.json', '../dataset/')
except AssertionError:
    pass

MetadataCatalog.get('coco_trash_train').thing_classes = ["General trash", "Paper", "Paper pack", "Metal", 
                                                         "Glass", "Plastic", "Styrofoam", "Plastic bag", "Battery", "Clothing"]

register_coco_instances함수를 통해, 미리준비된 coco format data를 등록시켜 주는 코드

만약 추가적인 처리가 필요하다면 load_coco_json함수를 이용해보자

마지막 get함수를 이용하여 trainset의 class들을 정의한다

 

 

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file('COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml'))

config 파일을 가져오는 코드

2번째 라인은 COCO-Detection dir에 있는 해당 yaml형식의 파일을 default cfg에 merge 해준다

* detectron2의 config 파일은 .yaml 타입으로 정의된다

 

 

 

cfg.DATASETS.TRAIN = ('coco_trash_train',)
cfg.DATASETS.TEST = ('coco_trash_test',)

cfg.DATALOADER.NUM_WOREKRS = 2

cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url('COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml')

cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.MAX_ITER = 15000
cfg.SOLVER.STEPS = (8000,12000)
cfg.SOLVER.GAMMA = 0.005
cfg.SOLVER.CHECKPOINT_PERIOD = 3000

cfg.OUTPUT_DIR = './output'

cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 10

cfg.TEST.EVAL_PERIOD = 3000

출처 : https://detectron2.readthedocs.io/en/latest/modules/solver.html

SOLVER 함수를 통해 하이퍼 파라미터들을 재정의 해준다. 만약 재정의를 하지 않는 파라미터들은 원 모델의 디폴트 파라미터로 설정된다.

 

 

 

 

def MyMapper(dataset_dict):
    dataset_dict = copy.deepcopy(dataset_dict)
    image = utils.read_image(dataset_dict['file_name'], format='BGR')
    
    transform_list = [
        T.RandomFlip(prob=0.5, horizontal=False, vertical=True),
        T.RandomBrightness(0.8, 1.8),
        T.RandomContrast(0.6, 1.3)
    ]
    
    image, transforms = T.apply_transform_gens(transform_list, image)
    
    dataset_dict['image'] = torch.as_tensor(image.transpose(2,0,1).astype('float32'))
    
    annos = [
        utils.transform_instance_annotations(obj, transforms, image.shape[:2])
        for obj in dataset_dict.pop('annotations')
        if obj.get('iscrowd', 0) == 0
    ]
    
    instances = utils.annotations_to_instances(annos, image.shape[:2])
    dataset_dict['instances'] = utils.filter_empty_instances(instances)
    
    return dataset_dict

mapper - input data를 어떤 형식으로 return할지 (따라서 augmnentation 등 데이터 전처리 포함 됨), 내가 정의해야됨 MMDetection은 config에 augmentation이 구현되있어서 따로 정의가 필요없지만, Detectron2는 내가 직접 정의해줘야된다

transform 과정은 여러분들도 쉽게 따라올 수 있는 코드이다. 그래서 그 이후 코드에 대해서 설명하려한다.

 

dataset_dict['image'] 를 알맞게 텐서로 변환시키고 + transpose로 차원도 맞춰준다.

 

annos부분은 내장되어있는 annoatations을 이용한다는 코드 같은데 정확한 것은 더 찾아봐야 될것 같다

 

 

 

class MyTrainer(DefaultTrainer):
    
    #재정의 과정 MyMapper을 mapper에 할당
    @classmethod
    def build_train_loader(cls, cfg, sampler=None):
        return build_detection_train_loader(
        cfg, mapper = MyMapper, sampler = sampler
        )
    
    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        if output_folder is None:
            os.makedirs('./output_eval', exist_ok = True)
            output_folder = './output_eval'
            
        return COCOEvaluator(dataset_name, cfg, False, output_folder)

재정의를 위한 함수

왜 재정의를 해야되냐?

>>> DefaultTrainer을 이용하여 build_train_loader 메서드를 재정의하여, custom dataloader을 만들 수 있기 때문

 

 

 

os.makedirs(cfg.OUTPUT_DIR, exist_ok = True)

trainer = MyTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

train 시키는 코드! 별 어려운 것은 없다고 본다