mmpretrain (classification) 사용법

mmpretrain 프레임워크는 리뉴얼 되면서 mmclassification과 mmSelfsUp (self-supervised learning) 이 통합되었다. task마다 docker를 만들어야 되다 보니 두 프레임워크가 합쳐진건 환영할만한 일이다.👏

 

 

아래 설명은 mmpretrain tag 1.0.0rc7 버전이 기준이며,

mmpretrain도 mmclassification 사용 방법과 크게 다르진 않다.

 

 

config를 구성하고 학습하면 되는데, 중요한 가지는 1) dataset 구성 2) config 구성 이 정도인 것 같다.

예시로 ImageNet-1k를 resnet50으로 학습해보는 과정을 진행해보겠다.

 

 

사람마다 다를 것 같긴 한데 나는 mmpretrain/model_config 라는 나의 config 폴더를 따로 만들어서 코드를 모두 붙여두고 한 방에 사용하는 편이다.

예를 들면, resnet50_8xb32_in1k.py 라는 config가 있는데, 코드는 이렇게 생겼다. 이게 전부이다.

# ../_base_ 폴더의 config 파일들을 끌어서 정의해놓음
_base_ = [
    '../_base_/models/resnet50.py', '../_base_/datasets/imagenet_bs32.py',
    '../_base_/schedules/imagenet_bs256.py', '../_base_/default_runtime.py'
]

 

이렇게 정의가 되어 있어서 하이퍼파라미터 튜닝을 할 경우 조금 불편한 부분이 발생한다. 저 네 개의 경로를 다 들어가서 수정을 하면 원래 코드에도 영향이 가기 때문에 ../_base_/models 이런 경로들을 모두 들어가서 나의 세팅 파일을 만들고 끌어서 쓰는 코드를 수정하고… 하는 것 보단 그냥 죄다 한 config 파일에 붙여서 내 입맛대로 수정하고 바로 실험해보는 편이다.

이 부분을 감안하고 보면 좋을 것 같다.

 

 

 

0. 환경 구성

• mmpretrain tag 1.0.0rc7 버전을 clone 해준다.

git clone -b v1.0.0rc7 https://github.com/open-mmlab/mmpretrain.git

• mm 프레임워크 시리즈 같은 경우는 Dockerfile이 있어서 도커로 환경 구성을 해주면 편하다.

 

1. Imagenet-1k 데이터 준비
다운로드는 https://www.kaggle.com/competitions/imagenet-object-localization-challenge/data 여기서 받을 수 있다.

 

 

2. Imagenet-1k 데이터 구조 확인
다운로드를 받게 되면 처음의 구조는 이런 형식이다. 우리가 사용할 데이터 부분만 보겠다.

imagenet/ILSVRC/Data/CLS-LOC # annotation을 만들 폴더는 여기
imagenet/ILSVRC/Data/CLS-LOC/train
imagenet/ILSVRC/Data/CLS-LOC/train/이미지 폴더
imagenet/ILSVRC/Data/CLS-LOC/train/이미지 폴더/이미지.JPEG # 1,283,166장
imagenet/ILSVRC/Data/CLS-LOC/val
imagenet/ILSVRC/Data/CLS-LOC/val/이미지.JPEG # 50,000장
imagenet/ILSVRC/Data/CLS-LOC/test # 학습에 사용하지 않음

 

 

3. 학습 annotation 파일 준비

• docs/en/advanced_guides/datasets.md 문서를 보면 txt 파일을 하기와 같은 방법으로 구성하라고 되어 있다.
• 이 형식에 맞게 구성해주면 된다. 코드를 첨부하지만 공부 겸 하고 싶으면 코드를 보지 않고 직접 구성해도 된다.

#이미지상대경로 #클래스
00000.png 2
00001.png 0
00002.png 33

from glob import glob
import os
from tqdm import tqdm
import xml.etree.ElementTree as elemTree

def main(base_dir):
    base_dir = os.path.join(base_dir, 'Data', 'CLS-LOC')
    phases = ['train', 'val']

    mem = []
    for phase in phases:
        save_txt_path = os.path.join(base_dir, f'{phase}.txt')
        if phase == 'train':
            image_dirs = sorted(glob(os.path.join(base_dir, phase, '**')))
            
            output = str()
            for i, image_dir in tqdm(enumerate(image_dirs)):
                image_paths = sorted(glob(os.path.join(image_dir, '*.JPEG')))
                
                for image_path in image_paths:
                    output += f"{'/'.join(image_path.split('/')[-2:])} {i}\\n"

                mem.append(image_dir.split('/')[-1])

            with open(save_txt_path, 'w') as f:
                f.writelines(output)
        
        else:
            output2 = str()
            image_dirs = sorted(glob(os.path.join(base_dir, phase, '**')))
            image_paths = glob(os.path.join(base_dir, phase, '*.JPEG'))
            
            for image_path in image_paths:
                xml_path = image_path.replace('Data', 'Annotations').replace('.JPEG', '.xml')
                tree = elemTree.parse(xml_path)
                class_name = tree.getroot()[5][0].text
                output2 += f"{image_path.split('/')[-1]} {mem.index(class_name)}\\n"
            
            with open(save_txt_path, 'w') as f:
                f.writelines(output2)

if __name__ == '__main__':
    base_dir = '/root/workspace/imagenet/ILSVRC'
    main(base_dir)

 

 

4. config 구성

• 위의 config/_base_ 폴더에 있는 config 코드들을 한 곳에 다 끌어모은 config 이다.
• Imagenet 학습 그대로 하는 것이기 때문에 하이퍼 파라미터 수정은 하지 않는다.
• 데이터셋 세팅이 중요하다.
• 글자를 자세히 읽어보면 뭔지 다 알 수 있게 쉽게 설명되어 있다. 여기서는 train, val 데이터로 학습을 진행하기 때문에 test dataloader 구성은 생략했다.

# dataset settings
dataset_type = 'ImageNet'
data_preprocessor = dict(
    num_classes=1000,
    # RGB format normalization parameters
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    # convert image from BGR to RGB
    to_rgb=True,
)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', scale=224),
    dict(type='RandomFlip', prob=0.5, direction='horizontal'),
    dict(type='PackInputs'),
]

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='ResizeEdge', scale=256, edge='short'),
    dict(type='CenterCrop', crop_size=224),
    dict(type='PackInputs'),
]

train_dataloader = dict(
    batch_size=64,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='imagenet/ILSVRC/Data/CLS-LOC/',
        split='train',
        ann_file='train.txt',
        pipeline=train_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=True),
)

val_dataloader = dict(
    batch_size=64,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='imagenet/ILSVRC/Data/CLS-LOC/',
        split='val',
        ann_file='val.txt',
        pipeline=test_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=False),
)
val_evaluator = dict(type='Accuracy', topk=(1, 5))

# If you want standard test, please manually configure the test dataset
test_dataloader = val_dataloader
test_evaluator = val_evaluator


# model settings
model = dict(
    type='ImageClassifier',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=2048,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
    ))

# defaults to use registries in mmpretrain
default_scope = 'mmpretrain'

# configure default hooks
default_hooks = dict(
    # record the time of every iteration.
    timer=dict(type='IterTimerHook'),

    # print log every 100 iterations.
    logger=dict(type='LoggerHook', interval=100),

    # enable the parameter scheduler.
    param_scheduler=dict(type='ParamSchedulerHook'),

    # save checkpoint per epoch.
    checkpoint=dict(type='CheckpointHook', interval=1),

    # set sampler seed in distributed evrionment.
    sampler_seed=dict(type='DistSamplerSeedHook'),

    # validation results visualization, set True to enable it.
    visualization=dict(type='VisualizationHook', enable=False),
)

# configure environment
env_cfg = dict(
    # whether to enable cudnn benchmark
    cudnn_benchmark=False,

    # set multi process parameters
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),

    # set distributed parameters
    dist_cfg=dict(backend='nccl'),
)

# set visualizer
vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(type='UniversalVisualizer', vis_backends=vis_backends)

# set log level
log_level = 'INFO'

# load from which checkpoint
load_from = None

# whether to resume training from the loaded checkpoint
resume = False

# Defaults to use random seed and disable `deterministic`
randomness = dict(seed=None, deterministic=False)


# optimizer
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001))

# learning policy
param_scheduler = dict(
    type='MultiStepLR', by_epoch=True, milestones=[30, 60, 90], gamma=0.1)

# train, val, test setting
train_cfg = dict(by_epoch=True, max_epochs=100, val_interval=1)
val_cfg = dict()
test_cfg = dict()

# NOTE: `auto_scale_lr` is for automatically scaling LR,
# based on the actual training batch size.
auto_scale_lr = dict(base_batch_size=256)

 

 

5. 학습하기

# multi-gpu
bash ./tools/dist_train.sh [config 파일 경로] [사용 GPU 수]

# single-gpu
python ./tools/train.py [config 파일 경로]

 

 

 

도움이 되셨다면 아무 광고나 클릭 한 번 부탁드립니다👍