MMCV Segmentation 입문기 2

안녕하세요. 지난 포스팅의 MMCV Segmentation 입문기 1에서는 MMCV Segmentation 라이브러리를 사용하기 위한 설치 과정과 제대로 설치되었는 지 확인해보는 시간을 가졌습니다. 오늘은 MMCV Segmentation을 사용하기 위한 기본적인 지식을 배워보도록 하겠습니다. 

 

1. Training Model Example

 

기본적으로 MMCV Segmentation을 이용한 학습은 tools 폴더에 있는 dist_train.sh 파일을 이용해야합니다. 이와 함께 넘겨줘야하는 것은 가장 중요한 configuration 파일이죠. MMCV Segmentation은 이 configuration 파일로 모든 설정을 맞추어줄 수 있으니 이후에 더 자세히 알아보도록 하겠습니다. 기본적인 학습 코드는 아래와 같습니다. 

 

- Single GPU

sh tools/dist_train.sh ${CONFIG_FILE} 1 [optional arguments]

 

- Multiple GPU

sh tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments]

 

여기서 [optional arguments]는 다음과 같습니다. 

 

--work-dir ${WORK_DIR} : 작업 디렉토리 정의

--resume-from ${CHECHPOINT_FILE} : 이전 체크포인트 파일에서 재시작 (학습 중단 후 재시작 할 때 사용)

--load-from ${CHECKPOINT_FILE} : 체크포인트 파일에 저장된 가중치만 로딩한 뒤 0 iteration부터 재학습 (미세 조정)

--deterministic : 훈련속도는 감소하지만 재현성을 확보

 

# checkpoints and logs saved in WORK_DIR=work_dirs/pspnet_r50-d8_512x512_80k_ade20k/
# If work_dir is not set, it will be generated automatically.
sh tools/dist_train.sh configs/pspnet/pspnet_r50-d8_512x512_80k_ade20k.py 8 --work_dir work_dirs/pspnet_r50-d8_512x512_80k_ade20k/ --deterministic

 

예를 들어 위 명령어를 실행하게 되면 configuration 파일은 pspnet_r50-d8_512x512_80k_ade20k.py에서 불러오게 됩니다. 그리고 GPU의 개수는 8개를 사용하게 되죠. 학습 시 체크포인트 파일은 --work_dir에 명시된 work_dirs/pspnet_r50-d8_512x512_80k_ade20k/에 저장됩니다. 이때, --deterministic을 사용하였기 때문에 재현성은 보장되죠. 

 

여기까지가 기본적인 학습 방법입니다. 이제부터 저희가 해야하는 것은 configuration 파일을 어떻게 명시할 것인지, 그리고 체크포인트 파일은 어디에 저장할 지만 결정하면 됩니다. 아주 간단하죠? 

 

2. Configuration 파일

 

일단, configuration 파일 이름부터 보도록 하죠.

 

pspnet_r50-d8_512x512_80k_ade20k.py

 

와... 뭔지는 정확하게 모르겠지만 되게 복잡해보입니다. 사실 MCCV 라이브러리에서는 모든 configuration 파일에는 네이밍 규칙이 존재합니다. 

 

{model}_{backbone}_[misc]_[gpu x batch_per_gpu]_{resolution}_{iterations}_{dataset}

 

위 형식이 configuration 파일의 네이밍 규칙입니다. 

 

{model} = pspnet : 사용할 segmentation 모델로 위의 경우에는 pspnet 사용

{backbone} = r50 : segmentation 모델의 backbone 네트워크로 위의 경우에는 resnet50 사용

[misc] : optional settings

[gpu x batch_per_gpu] : 사용한 GPU 개수와 GPU 당 배치 사이즈

{resolution} = 512 x 512 : 학습한 데이터셋의 해상도로 위의 경우에는 512 x 512 크기의 영상 사용

{iterations} = 80k : 학습 횟수로 위의 경우에는 80k 학습

{dataset} = ade20k : 사용한 데이터셋으로 위의 경우에는 ade20k라는 데이터셋 사용

 

이렇게 보니 생각보다 어렵지 않습니다. 이렇게 configuration 파일을 만든 뒤 sh 명령어를 이용해서 학습을 시작하면 되겠네요! 그렇다면 해당 py 파일이 어떻게 생겼는 지 뜯어서 확인해보도록 하죠.

 

_base_ = [
    '../../_base_/models/segformer.py',
    '../../_base_/datasets/ade20k_repeat.py',
    '../../_base_/default_runtime.py',
    '../../_base_/schedules/schedule_160k_adamw.py'
]

# model settings
norm_cfg = dict(type='SyncBN', requires_grad=True)
find_unused_parameters = True
model = dict(
    type='EncoderDecoder',
    pretrained='pretrained/mit_b2.pth',
    backbone=dict(
        type='mit_b2',
        style='pytorch'),
    decode_head=dict(
        type='SegFormerHead',
        # type='MLPHead',
        in_channels=[64, 128, 320, 512],
        in_index=[0, 1, 2, 3],
        feature_strides=[4, 8, 16, 32],
        channels=128,
        dropout_ratio=0.1,
        num_classes=150,
        norm_cfg=norm_cfg,
        align_corners=False,
        decoder_params=dict(embed_dim=768),
        loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

# optimizer
optimizer = dict(_delete_=True, type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.01,
                 paramwise_cfg=dict(custom_keys={'pos_block': dict(decay_mult=0.),
                                                 'norm': dict(decay_mult=0.),
                                                 'head': dict(lr_mult=10.)
                                                 }))

lr_config = dict(_delete_=True, policy='poly',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 power=1.0, min_lr=0.0, by_epoch=False)

data = dict(samples_per_gpu=2)
evaluation = dict(interval=16000, metric='mIoU')

 

방금 위의 코드를 보셨다면 매우 당황하셨을 겁니다. 너무 긴 거 같은데?... 이걸 어떻게 쓴다는 걸까요? 한번 자세히 뜯어보도록 하죠. 기본적으로 MCCV의 configuration.py 파일의 구성목록은 model, data, optimization, evaluation으로 이루어집니다. 이 정도만 알아도 학습을 수행하는 큰 문제가 없으니 중요한 몇몇 configuration만 보도록 하죠. 

 

- Model Configuration

 

norm_cfg = dict(type='SyncBN', requires_grad=True)
find_unused_parameters = True
model = dict(
    type='EncoderDecoder',
    pretrained='pretrained/mit_b2.pth',
    backbone=dict(
        type='mit_b2',
        style='pytorch'),
    decode_head=dict(
        type='SegFormerHead',
        # type='MLPHead',
        in_channels=[64, 128, 320, 512],
        in_index=[0, 1, 2, 3],
        feature_strides=[4, 8, 16, 32],
        channels=128,
        dropout_ratio=0.1,
        num_classes=150,
        norm_cfg=norm_cfg,
        align_corners=False,
        decoder_params=dict(embed_dim=768),
        loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

 

사용할 segmentation 모델의 configuration이 정의된 부분입니다. 제가 사용할 segmentation 모델은 SegFormer입니다. 기본적으로 딕셔너리 자료형으로 이루어져 있기 때문에 생각보다 쉽게 이해하실 수 있을겁니다. 자, 여기서 model의 type은 EncoderDecoder로 왠만하면 고정이 되는 것 같습니다. 그리고 pretrained는 pretrained된 backbone 네트워크의 경로를 적어주면 됩니다. 

 

다음으로는 bacbone을 정의해줍니다. 본 논문에서는 mit_b2라고 백본 네트워크를 정의하였습니다. 실제로는 transformer encoder를 사용했는데... 이 부분은 명확하게 이해는 되지 않습니다만, 만약 다른 종류의 encoder를 사용하고 싶으시다면 mmseg/models/backbones에 더 많은 종류의 백본이 존재하기 때문에 한번 활용해보시길 바랍니다. 그리고 style을 pytorch로 정의하게 되는 데 이 부분도 딱히 건들일은 없겠습니다. 

 

이제 decoder 부분을 정의해보도록 하겠습니다. 디코더의 타입은 해당 논문에서는 SegFormerHead라고 작성하였는데 다른 종류의 decoder를 사용하고 싶으시면 mmseg/models/decoder_heads에 더 많은 종류의 디코더가 존재하기 때문에 확인해보시길 바랍니다. 여기에 추가적으로 디코더와 관련된 다양한 파라미터들(각 계층의 채널 개수, 출력 계층의 최종 아웃풋 채널 개수, ...)을 넘겨줍니다. 또한, loss_decode에서는 예측과 ground truth 사이의 손실함수를 계산하게 되죠. 

 

마지막으로 train_cfg와 test_cfg 역시 위 방법으로 고정하게 학습 및 평가가 이루어지기 때문에 특별히 건들일은 없는 거 같습니다. 

 

- Optimizer Configuration & Learning rate Scheduler Configuration

 

# optimizer
optimizer = dict(_delete_=True, 
                 type='AdamW', 
                 lr=0.00006, 
                 betas=(0.9, 0.999), 
                 weight_decay=0.01,
                 paramwise_cfg=dict(custom_keys={'pos_block': dict(decay_mult=0.),
                                                 'norm': dict(decay_mult=0.),
                                                 'head': dict(lr_mult=10.)
                                                 }))
                                                 
lr_config = dict(_delete_=True, policy='poly',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 power=1.0, min_lr=0.0, by_epoch=False)

 

최적화 함수는 모델에 비해 비교적 간단합니다. type에는 사용할 최적화 알고리즘의 종류를 선택하시면 됩니다. 그리고 나머지에는 최적화 함수에 필요한 파라미터들 (lr, betas, weight_decay, ...)를 넘겨주면 됩니다. 이 부분은 pytorch를 사용하실 때 optimizer를 정의하는 부분과 크게 다르지는 않습니다. 

 

다음으로 학습률을 어떤식으로 변경할 것인가에 대한 스케쥴러 configuration도 함께 넘겨주면 됩니다. 

 

- Evaluation Configuration

 

evaluation = dict(interval=16000, metric='mIoU')

 

해당 configuration은 몇 번의 interval 마다 평가를 할 것인지, 그리고 어떤 평가방법을 선택할 지에 대한 파라미터만 넘겨주면 됩니다.