[BoostCamp AI Tech / PyTorch] Day6 - PyTorch 프로젝트 구조 이해하기

PyTorch : PyTorch 프로젝트 구조 이해하기

---

Introduction

• 많이들 Jupyter notebook을 통해 ML/DL 입문을 하고 코드를 접한다.
• 하지만 Jupyter notebook에는 한계가 존재한다.
  • 배포 및 공유의 어려움 $\rightarrow$ 재현의 어려움과 실행순서의 꼬임
• 이를 해결하고자 Deep Learning 코드를 하나이 프로그램으로 작성할 필요가 있다.
  • 개발의 용이성 확보, 유지보수 향상
• 앞서 Python의 OOP 개념을 배운 것이 여기써 쓰인다.

PyTorch Template

• 파이토치의 템플릿의 종류는 다양한데 그 중 일부를 소개하겠다.
• 사용자의 필요에 따라 수정하여 사용한다.
• train, data_loader, model, config, logging, metric, utils 등 다양한 모듈로 분리하여 프로젝트를 사용한다.
• 실제로 많은 기업들과 대학원 연구실에서 템플릿을 활용하는 방식을 사용한다.

PyTorch Template Project

pytorch-template/
│
├── train.py - main script to start training
├── test.py - evaluation of trained model
│
├── config.json - holds configuration for training
├── parse_config.py - class to handle config file and cli options
│
├── new_project.py - initialize new project with template files
│
├── base/ - abstract base classes
│   ├── base_data_loader.py
│   ├── base_model.py
│   └── base_trainer.py
│
├── data_loader/ - anything about data loading goes here
│   └── data_loaders.py
│
├── data/ - default directory for storing input data
│
├── model/ - models, losses, and metrics
│   ├── model.py
│   ├── metric.py
│   └── loss.py
│
├── saved/
│   ├── models/ - trained models are saved here
│   └── log/ - default logdir for tensorboard and logging output
│
├── trainer/ - trainers
│   └── trainer.py
│
├── logger/ - module for tensorboard visualization and logging
│   ├── visualization.py
│   ├── logger.py
│   └── logger_config.json
│  
└── utils/ - small utility functions
    ├── util.py
    └── ...

• PyTorch Template Project는 아마도 많이 사용하는 형태의 템플릿이 아닐까 싶다.
• train.py, test.py
  • train 및 test가 실제로 실행되는 코드이다.
  • 템플릿의 핵심이 되는 코드들 중 하나이다.
• config.json, parse_config.py
  • config.json은 모델의 학습 또는 데이터 관련 등 ML 프로젝트의 설정을 담고 있는 파일이다.
  • parse_config.py는 ConfigParser라는 클래스가 작성되어있다.

      @classmethod
      def from_args(cls, args, options=''):
          """
          Initialize this class from some cli arguments. Used in train, test.
          """
          for opt in options:
              args.add_argument(*opt.flags, default=None, type=opt.type)
          if not isinstance(args, tuple):
              args = args.parse_args()
    
          if args.device is not None:
              os.environ["CUDA_VISIBLE_DEVICES"] = args.device
          if args.resume is not None:
              resume = Path(args.resume)
              cfg_fname = resume.parent / 'config.json'
          else:
              msg_no_cfg = "Configuration file need to be specified. Add '-c config.json', for example."
              assert args.config is not None, msg_no_cfg
              resume = None
              cfg_fname = Path(args.config)
                
          config = read_json(cfg_fname)
          if args.config and resume:
              # update new config for fine-tuning
              config.update(read_json(args.config))
    
          # parse custom cli options into dictionary
          modification = {opt.target : getattr(args, _get_opt_name(opt.flags)) for opt in options}
          return cls(config, resume, modification)
    

    핵심적인 코드 중 하나인데 입력으로 들어온 argument들을 파싱하여 ConfigParser 객체로 반환하는 역할을 한다.
    여기에 같이 구현된 mangling 함수 중 __getitme__이 있는데

      def __getitem__(self, name):
          """Access items like ordinary dict."""
          return self.config[name]
    

    이 함수는 ConfigParser 객체에 한해서 []로 딕셔너리 접근을 가능하게 만들어주는 역할을 한다.

• data_loader/, data/
  • 이 폴더에는 data_loader.py와 학습을 위한 데이터가 들어있다.
  • data_loader는 모델 학습에서 중요하게 쓰이는데 이 부분의 구현이 된 부분이다.
• model/
  • 이 폴더는 우리가 학습에 사용할 모델의 구조가 구현된 부분이다. 물론 모델만 있는 것은 아니지만 대부분 model.py에 새로운 모델을 클래스로 구현해서 train.py나 test.py에 적용하여 모델을 생성한다.
• trainer/
  • 가장 핵심이 되는 파일이다.
  • model.py에 존재하는 모델을 사용해서 학습을 본격적으로 진행하는 코드이다.
  • train.py에 있는 main 함수에 다음과 같은 코드가 있다.

      def main(config):
          ...
          trainer = Trainer(model, criterion, metrics, optimizer,
                              config=config,
                              device=device,
                              data_loader=data_loader,
                              valid_data_loader=valid_data_loader,
                              lr_scheduler=lr_scheduler)
    
          trainer.train()
    

    여기서 trainer가 trainer.py에 구현된 Trainer클래스의 객체를 사용한다.

• utils/
  • 프로젝트에 사용되는 다양한 짜잘짜잘한 함수들을 구현한 파일이 있다.

PyTorch Lightning

PyTorch Lightning

• PyTorch Lightning은 간단하게 말하면 TensorFlow의 Keras 같은 존재이다.
• PyTorch의 nn.Module을 상속해서 모델을 생성하는 방식을 많이 사용하지만 Multi-GPU나 TPU, 분산학습 등 복잡한 실험환경에서는 코드가 길어지는 경우가 발생한다.
  이를 해결하고자 좀 더 간결하고 집중하고자 하는 곳에 더 집중할 수 있게 도움을 주는 것이 PyTorch Lightning의 역할이다.

• 아마 sci-kit learn에서 model을 만들고 model.fit으로 한방에 학습을 시켜본 기억이 있을 것이다.
  PyTorch Lightning은 이렇게 코드 자체를 획기적으로 줄여주는 역할을 한다.

# ----------------
# TRAINING LOOP
# ----------------
num_epochs = 1
for epoch in range(num_epochs):

  # TRAINING LOOP
  for train_batch in mnist_train:
    x, y = train_batch

    logits = pytorch_model(x)
    loss = cross_entropy_loss(logits, y)
    print('train loss: ', loss.item())

    loss.backward()

    optimizer.step()
    optimizer.zero_grad()

  # VALIDATION LOOP
  with torch.no_grad():
    val_loss = []
    for val_batch in mnist_val:
      x, y = val_batch
      logits = pytorch_model(x)
      val_loss.append(cross_entropy_loss(logits, y).item())

    val_loss = torch.mean(torch.tensor(val_loss))
    print('val_loss: ', val_loss.item())

이랬던 코드를

# train
model = LightningMNISTClassifier()
trainer = pl.Trainer()

trainer.fit(model)

이렇게 만들 수 있다.

PyTorch Lightning 써보기

• 참고링크

• Model 구현

class Classifier(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 10)
        )
    
    def forward(self, x):
        pass

    def training_step(self, batch, batch_idx):
        pass

    def validation_step(self, batch, batch_idx):
        pass

    def test_step(self, batch, batch_idx):
        pass
    
    def configure_optimizers(self):
        pass

• PyTorch Lightning을 활용한 모델은 위와 같은 구성으로 구현하면된다.
• training_step과 configure_optimizers는 반드시 구현해야한다.
• def __init__(self)
  • 이 부분은 모델 그래프를 구현하는 곳이다.
• def forward()

    def forward(self, x):
        return self.model(x)
  

  모델의 추론 결과를 제공할 때 사용한다. nn.Module처럼 반드시 정의하는 부분은 아니지만 다른 메서드 구현때 편리하므로 구현하는 것이 좋다.

• def training_step()

     def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        return loss
  

  training_step은 학습 루프의 body이다. 이 메소드에는 argumetn에 train dataloader가 제동하는 batch와 batch index가 주어지고 이를 활용해 학습 loss를 계산한다.
  PyTorch Lightning은 CPU, GPU 세팅을 따로 설정하지 않아도 trainer 설정에 맞춰 자동으로 typecasting을 해준다.

• configure_optimizers

    def configure_optimizers(self):
        optimizer = Adam(self.parameters())
        return optimizer
  

  configure_optimizers는 모델의 최적 파라미터를 찾을 때 사용할 optimizer와 scheduler를 구현한다.
  학습모델이 여러개면 리스트로 리턴을 해주면된다. 여기서는 모델이 1개이므로 Adam만 사용했다.

• 학습 및 테스트

pl.seed_everything(args.seed) # 재생산을 위한 랜덤시드 고정

# dataloaders
dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())
train_dataset, val_dataset = random_split(dataset, [55000, 5000])
test_dataset = MNIST('', train=False, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=args.batch_size)
val_loader = DataLoader(val_dataset, batch_size=args.batch_size)
test_loader = DataLoader(test_dataset, batch_size=args.batch_size)

# model
model = Classifier()

# training
trainer = pl.Trainer(max_epochs=args.n_epochs, gpus=args.n_gpus)
trainer.fit(model, train_loader, val_loader)

trainer.test(test_dataloaders=test_loader)

• Keras에서 MNIST를 학습하고 테스트할 때와 유사한 코드를 보여주고 있다.
• 위에서 모델을 구현해두면 간단하게 모델을 학습, 테스트 할 수 있는 것을 볼 수 있다.