[PyTorch] Pytorch template을 이용한 pytorch 프로젝트

jypter는 python코드나 모델의 실행결과를 편하고 쉽게 확인할 수 있다. 하지만, jupyter만으로는 모델 학습, 평가 및 log를 기록하는 것이 힘들고, 다른 사람들이 모델을 사용할 수 있도록 배포를 하는 것이 힘들다. 따라서, jupyter는 개발 초기 단계에서 학습 과정을 확인하거나, 실험, 디버깅 등에서만 사용하는 것이 좋다.

모델을 배포하고, 개발의 용이성을 위해서는 OOP와 모듈을 활용한 하나의 프로젝트로 관리하는 것이 좋다. 이러한 pytorch 프로젝트의 개발을 위해, github등에서는 여러 template을 찾을 수 있다. 이 포스트에서는 여러 template중, victoresque님의 pytorch-template를 사용해 볼 예정이다.

사용 방법

1. 해당 github repository에 접속한다(https://github.com/victoresque/pytorch-template)
2. 우측 상단의 Code를 누른 뒤 나온 clone URL을 통해 저장소를 clone한다.

git clone https://github.com/victoresque/pytorch-template.git

3. clone한 폴더에서 new_project.py를 실행시켜 기본 template으로 구성된 새로운 프로젝트 폴더를 만든다.

python new_project.py ../NewProject

4. 만들어진 프로젝트 폴더에서 모델, DataLoader, loss 등을 정의하고, 설정에서 hyperparameter등을 설정한 뒤, train,py와 test.py로 모델을 학습 및 테스트한다.

프로젝트 구조

프로젝트의 구조는 기본적으로 다음과 같다.

pytorch-template/
│
├── train.py - training 시작
├── test.py - 훈련된 모델 평가
│
├── config.json - 훈련을 위한 설정 파일
├── parse_config.py - 설정 파일과 cli 옵션을 parse하기 위한 클래스
│
├── new_project.py - template파일들로 새로운 프로젝트 폴더 생성
│
├── base/ - custom dataloader, 모델, trainer 생성을 위한 abstract base classes
│   ├── base_data_loader.py
│   ├── base_model.py
│   └── base_trainer.py
│
├── data_loader/ - 데이터 관련 처리 및 dataloader용 파일
│   └── data_loaders.py
│
├── data/ - 입력 데이터가 저장 될 폴더
│
├── model/ - 모델, loss, metric을 저장할 폴더
│   ├── model.py
│   ├── metric.py
│   └── loss.py
│
├── saved/
│   ├── models/ - 훈련된 모델 및 checkpoint를 저장
│   └── log/ - tensorboard를 위한 logdir 및 log를 저장
│
├── trainer/ - trainers
│   └── trainer.py
│
├── logger/ - tensorboard 시각화와 logging을 위한 모듈
│   ├── visualization.py
│   ├── logger.py
│   └── logger_config.json
│  
└── utils/ - 유틸리티 function
    ├── util.py
    └── ...

모델 학습과 평가는 train.py와 test.py의 실행을 통해 이루어진다.

config.json파일에는 프로젝트 이름, gpu 수를 지정하고, 모델, dataloader, optimizer 등 학습에 필요한 객체들의 이름을 지정하여 코드를 수정하지 않고도 원하는 모델과 data, loss등을 사용할 수 있다.

data_loader/data_loaders.py에는 모델에 데이터를 공급할 DataLoader객체가 있다. pytorch의 DataLoader를 상속하거나, template에 있는 base/base_data_loader.py에 있는 BaseDataLoader객체를 상속받아 사용자가 원하는 dataloader를 만들 수 있다. BaseDataLoader사용 시, validation set 비율을 지정함으로써 쉽게 train/validation split을 할 수 있다.

model/ 폴더에는 모델, metric, loss를 저장할 수 있다. 모델은 dataloader와 마찬가지로 base_model.py의 BaseModel을 상속하여 만들 수 있다. 훈련에 사용될 loss와, 평가에 사용될 metric들도 각 모듈에 정의하여 사용할 수 있다.

train.py를 실행시켜 모델을 학습하게 되면, 실제 train/validation loop과 진행 상황 출력, logging등은 모두 trainer.py에서 일어난다. trainer.py가 잘 작성되어 있어 큰 수정 없이도 웬만한 모델은 학습시킬 수 있으며, 만약 일부 내용을 수정하거나 다른 logic이 필요하다면 base_trainer.py의 BaseTrainer를 상속받아 원하는 trainer를 만들 수도 있다.

utils/ 폴더에는 json파일 입출력, 단일 및 다중 gpu 지원, metric 기록을 위한 다양한 utility function이 정의되어있다.

직접 사용해보기

새로운 모델, loss, metric을 해당 template에 적용시켜 훈련한다고 가정한다.

여기서 사용할 모델은 음원 분리에 사용되는 딥러닝 모델 WaveUNet으로, 해당 모델과 관련 metric등을 이 포스트에서 설명하기에는 내용이 길어질 수 있어 추후 다른 포스트에서 정리하겠습니다.

우선, 사용하고자 하는 모델을 model/mode.py에 작성한다.
일반적인 pytorch 모델처럼 nn.Module을 상속하여 사용해도 되지만, BaseModel을 상속하면 내부에 __str__이 잘 정의되어 있어 모듈의 파라미터와 구조를 쉽게 출력할 수 있고, method도 forward만 구현해주면 된다.

class WaveUNet(BaseModel):
	 def __init__(self, n_level=12, n_source=4):
        super().__init__()
        self.level = n_level

        layers = [DownSampling(in_ch=1, out_ch=24, kernel_size=15)]
        # ...
	
    def forward(self, x: Tensor):
        layer_to_concat = []
        layer_to_concat.append(x)
        for layer in self.net[0 : self.level]:
            x = layer(x)
        # ...

다음으로, 학습에 사용될 loss function을 model/loss.py에 정의해준다. 여기서는 mse loss를 사용할 것이기 때문에, 간단하게 한 줄 정도로 만들 수 있다.

def mse_loss(output, target):
    return F.mse_loss(output, target)

다음으로, validation, test단계에서 모델의 평가지표로 사용되는 metric을 model/metric.py에 정의해준다.
보통 metric은 한 batch내의 metric을 계산한 뒤, 모든 batch의 metric을 더하고 batch 수로 나누는 식으로 평균을 계산하므로, metric의 parameter로 모델의 output과 target을 지정해주고, metric값만 return해주면 쉽게 사용가능하다.

def sdr(output, target):
    loss_fn = ModifiedSDR().to("cuda")
    y_target_acc = target[:, 0:3]
    y_target_acc = torch.flatten(y_target_acc, end_dim=1)
    y_pred_acc = output[:, 0:3]
    y_pred_acc = torch.flatten(y_pred_acc, end_dim=1)
    # ...

다음으로 dataloader를 data_loader/data_loaders.py에 정의한다. torchvision과 같이 사용할 수 있는 Dataset객체가 있다면 dataloader만 정의하면 되고, 만약 custom dataset이 필요하다면 dataset와 dataloader를 함께 만들어주면 된다.
dataloader는 BaseDataLoader를 상속받아 super().__init__()만 잘 호출해주면 된다.

class MUSDBDataset(Dataset):
    def __init__(self, data_dir: str, train=True):
        self.data_dir = os.path.normpath(data_dir)
        self.music_list = []
        # ...

class MUSDBDataloader(BaseDataLoader):
    def __init__(self, data_dir, batch_size, shuffle=True, validation_split=0.2, num_workers=1, training=True):
        self.data_dir = data_dir
        self.dataset = MUSDBDataset(self.data_dir, train=training)
        super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)

마지막으로 config.json에서 만들어둔 model, loss, metric, dataloader를 지정하고, 그 외에 optimizer, scheduler 및 각종 hyperparameter들을 설정해주면 된다.
설정 파일 내에서 객체의 이름들만 지정해주어도, parse_config.py에 의해 해당하는 이름의 객체를 만들어 반환해 주기 때문에 직접 코드를 수정하지 않고도 자유롭게 모델이나 dataloader등을 설정할 수 있다.
각 설정들에 대한 설명은 github repository의 README.md에서 확인할 수 있다. 설정은 아래 예시 외에도 원하는 것을 넣어줄 수 있다.

{
    "name": "WaveUNet",
    "n_gpu": 1,

    "arch": {
        "type": "WaveUNet",
        "args": {}
    },
    "data_loader": {
        "type": "MUSDBDataloader",
        "args":{
            "data_dir": "./data",
            "batch_size": 32,
            "shuffle": true,
            "validation_split": 0.2,
            "num_workers": 2
        }
    },
    "optimizer": {
        "type": "Adam",
        "args":{
            "lr": 0.0001,
            "weight_decay": 0,
            "amsgrad": true
        }
    },
    "loss": "mse_loss",
    "metrics": [
        "sdr"
    ],
    "lr_scheduler": {
        "type": "None",
        "args": "None"
    },
    "trainer": {
        "epochs": 5,

        "save_dir": "saved/",
        "save_period": 1,
        "verbosity": 2,
        
        "monitor": "min val_loss",
        "early_stop": 2,

        "tensorboard": false
    }
}

설정이 모두 끝났으면 cli에서 python train.py -c config.json를 입력해 학습을 시작할 수 있다.
trainer.py에 의해 웬만한 상황에서는 학습이 잘 진행되지만, 수정이 필요한 경우 trainer.py를 직접 수정하거나 BaseTrainer를 상속하여 자신만의 training, validation step과 logging등을 만들 수 있다.
앞서 정의한 metric인 SDR의 경우, 기존 trainer를 사용하면 metric이 잘못 계산되므로, _train_epoch와 _valid_epoch에서 metric을 저장하는 부분을 약간 수정하여 사용했다.

def _train_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Integer, current training epoch.
        :return: A log that contains average loss and metric in this epoch.
        """
        self.model.train()
        self.train_metrics.reset()
        # ...
        
            for met in self.metric_ftns:
            	# 원래는 self.train_metrics.update(met.__name__, met(output, target))만 존재
                sdr, count = met(output, target)
                self.train_metrics.update(met.__name__, sdr, count)

학습 시작 시, 모델의 구조와 파라미터 수를 출력해주고, 학습 진행 상황, train및 validation loss/metric, checkpoint 저장 위치, earlystopping 여부 등을 출력해준다.