[논문 구현] Neural Style Transfer In Pytorch

이전 글에서 포스팅 했던 A Neural Algorithm of Artistic Style 논문을 토대로 작업환경은 Colab을, Pytorch 오픈소스를 사용해 이미지 합성을 해보겠습니다.

논문을 먼저 보시고 싶은 분들은 아래 포스팅을 눌러주시면 됩니다.

[논문 이해하기] A Neural Algorithm of Artistic Style by Leon A. Gatys,  Alexander S. Ecker, Matthias Bethge

제가 소개해드릴 모델은 논문을 토대로 Pytorch를 사용해 기존의 VggNet(=pretrained model)을 Pretrained model로 불러와 훈련을 시키는 과정 속에서 loss값을 최소화해 합성된 이미지를 출력하는 모델입니다.

---

데이터를 저장할 폴터를 생성합니다.

# Create a folder to store data
!mkdir data

 

이미지 데이터 불러오기

저는 구글 드라이브에 있는 사진을 불러왔기 때문에 구글 드라이브와 연동시키는 함수 dirve.mount('path')를 사용했습니다.

from google.colab import drive
drive.mount('/content/drive')
# Mounted at /content/drive

연동이 완료가 되면 '사진이 구글 드라이브 내에 어떤 경로에 있는지 파악' 후 불러 옵니다.

잘 불러져 왔는지 사진을 확인해 봅니다.

# Check Content image and Style image
content_img
style_img

Transformation 적용하기

저는 이 과정에서 어려움을 겪었었는데요, 그것은 바로 CMYK 사진입니다. 결과 출력값을 보시다시피 torch.Size([3, 256, 384]) False 에서 3이라는 것이 채널이 3개가 있는 뜻입니다(RGB). 그러나 제가 load한 이미지는 CMYK의 채널이 4개였던 것이었죠. 그래서 RGB로 바꿔도 보고 이것저것 해보았었는데 그 결과, 그냥 RGB인 사진으로 잘 고르자 였습니다...
혹시 채널 개수가 4여서 오류가 뜨시는 분들은 가급적이면 RGB image를 load 해주시면 됩니다.

Transformation이 적용된 image를 시각화합니다.









그러면 이런식으로 왼쪽에는 Content Image, 오른쪽에는 Style Image가 들어갑니다.

Style transfer 구현하기

여기서 pretrained model인 VGG19 model를 불러왔습니다.
'cuda'라는 것은 GPU이름인데 본인 device가 gpu가 작동되는 환경이면 쓰고 아니면 cpu를 쓰겠다 이 뜻입니다.

#파라미터를 freeze 하는 이유를 알고 싶다면? 여기

 

---

 

손실함수 정의하기

gram matrix를 계산하는 함수입니다. gram matrix는 style loss를 계산할 때 필요합니다.
style loss는 input tensor와 style image의 gram matrix가 동일한 값을 갖는 방향으로 학습을 진행합니다.

 

• content loss

content feature와 input feature 사이의 MSE를 계산합니다.

 

• style loss

input과 style의 gram matrix가 동일해지도록 학습합니다.

pre-trained VGG19을 사용하여 content와 style image의 feature을 추출합니다.

추출한 feature를 확인합니다.

 

---

Optimizer 정의하기
input_tensor은 content image을 복사하여 생성합니다. optimizer은 input_tensor의 pixel value를 갱신합니다. 즉, 학습을 진행할수록 모델의 파라미터가 갱신되는 것이 아니라 input_tensor의 값이 갱신됩니다.

모든 준비는 끝났습니다.

---

실행하기

하이퍼 파마라미터를 정의합니다. 사진에는 300 epochs라고 나와있지만 저는 5000 까지도 해봤습니다.

---

epoch 5000이 되었을 때, 알게 되는 점이 있습니다. 일단 당연하게 style loss는 계속 떨어지고 있었고, content loss가 epoch 2200 때부터 5,0 값으로 유지하는 것을 알게되었습니다.


결과 시각화하기

이 합성된 출력 이미지를 구글 드라이브에 저장하는 방법입니다.

output_img.save('저장할 경로/저장할이름.png', 'png')


마지막으로, 많은 도움을 주신 @Inhovation97님께 감사합니다


전체 코드 주소입니다.

Neural Style Transfer for Pytorch.ipynb

출처 코드는 아래 github 주소입니다.

GitHub - PacktPublishing/PyTorch-Computer-Vision-Cookbook: PyTorch Computer Vision Cookbook, Published by Packt

참고 문헌
[1] 해당 논문