style transfer model

사진을 특정 화풍으로 변환하는 모델

pretrained vgg16 활용

 

전체코드

import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import matplotlib.pyplot as plt

# VGG 모델 준비 및 ReLU 설정 해제
vgg16 = models.vgg16(pretrained=True).features
for layer in vgg16:
    if isinstance(layer, nn.ReLU):
        layer.inplace = False

# Gram matrix 함수
def gram_matrix(features):
    with torch.no_grad():
        _, C, H, W = features.size()
        features = features.view(C, H * W)
        gram = torch.mm(features, features.t()) / (C * H * W)
    return gram

# 특징 추출 함수
def get_features(image, model, layer_idx):
    content_feature = None
    style_features = []
    x = image
    
    for i, layer in enumerate(model):
        x = layer(x)
        
        if i == layer_idx['content']:
            content_feature = x.clone()
        elif i in layer_idx['style']:
            style_features.append(gram_matrix(x))
    
    return content_feature, style_features

# 이미지 전처리 및 로드
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

content_image = transform(Image.open("your content image")).unsqueeze(0)
style_image = transform(Image.open("your style image")).unsqueeze(0)

# 레이어 인덱스 설정
layer_idx = {'content': 21, 'style': [0, 5, 10, 16, 26]}

# 콘텐츠 및 스타일 특징 추출
with torch.no_grad():
    content_features_from_content, style_features_from_content = get_features(content_image, vgg16, layer_idx)
    content_features_from_style, style_features_from_style = get_features(style_image, vgg16, layer_idx)

# 생성 이미지 초기화 (반복문 외부에서 한 번만 생성)
generated_image = content_image.clone().requires_grad_(True)

# 옵티마이저 설정
optimizer = torch.optim.Adam([generated_image], lr=0.0001)
criterion = nn.MSELoss()

content_weight = 0.3
style_weight = 1000

# 학습 루프
num_iterations = 1000
for i in range(num_iterations):
    optimizer.zero_grad()
    
    # 생성 이미지의 콘텐츠 및 스타일 특징 추출
    content_features_from_generated, style_features_from_generated = get_features(generated_image, vgg16, layer_idx)
    
    # 콘텐츠 손실 계산
    content_loss = criterion(content_features_from_content, content_features_from_generated)
    
    # 스타일 손실 계산
    style_loss = sum(criterion(style_gram, generated_gram)
                     for style_gram, generated_gram in zip(style_features_from_style, style_features_from_generated))
    
    # 총 손실 계산
    total_loss = content_weight * content_loss + style_weight * style_loss
    
    # 손실 값 출력
    print(f"Iteration {i+1}/{num_iterations}, Total Loss: {total_loss.item()}")
    
    # 역전파와 옵티마이저 업데이트
    total_loss.backward()
    optimizer.step()
    
    # 생성 이미지 값의 범위를 유지
    with torch.no_grad():
        generated_image.clamp_(0, 1)

    # 주기적으로 생성된 이미지 시각화
    if (i + 1) % 5 == 0 or i == num_iterations - 1:
        plt.figure()
        unloader = transforms.ToPILImage()
        image = generated_image.cpu().clone().squeeze(0)  # 차원 줄이기
        image = unloader(image)
        plt.imshow(image)
        plt.title(f"Generated Image at Iteration {i+1}")
        plt.axis('off')
        plt.show()

작성개요

1. 두 이미지를 준비

콘텐츠 이미지: 변환하고자 하는 원본 사진

스타일 이미지: 적용하고 싶은 화풍이 담긴 사진(수묵화, 르네상스화풍, 아크릴화, 수채화, 모네의 그림 등)

 

2. 콘텐츠 레이어와 스타일 레이어를 정하고 특징 추출함수 정의

CNN모델은 네트워크가 깊어질수록 더 고수준과 추상적인 정보를 학습하고 복잡한 이미지 정보를 잡아낸다. 초기 레이어는 엣지나 색상같은 저수준 특징을 학습하고, 깊은레이어로 가면 전체적인 구조나 스타일을 학습한다.

콘텐츠 레이어: CNN 모델의 중간~깊은 단계에 있는 하나의 convolution 레이어를 활용하여 이미지의 구도나 형태 보존

스타일 레이어: CNN 모델의 다양한 층에 있는 여러 개의 convolution 레이어를 활용하여 다양한 특징을 포착 

 

3. 손실함수 정의(MSE사용)

콘텐츠손실: 콘텐츠이미지와 생성된이미지의 콘텐츠레이어를 비교
스타일손실: 스타일이미지와 생성된이미지의 스타일레이어를, 각 레이어의 그람메트릭스 차이를 구해서 합함

총 손실: 콘텐츠가중치*콘텐츠손실 + 스타일가중치*스타일손실

 

4. 이미지 생성

 

원래 코드는 이보다 좀 더 간단했는데 역전파과정에서 자꾸만 오류가 생겨 이런저런 설정을 추가했다.

 

vgg16 = models.vgg16(pretrained=True).features
for layer in vgg16:
    if isinstance(layer, nn.ReLU):
        layer.inplace = False

모델을 불러오고, ReLU layer에 설정을 추가함

ReLU레이어에서 연산할 때 입력텐서의 메모리 위에 덮어쓰기를 하는데 역전파 과정에서 문제가 될 수 있어 덮어쓰기 설정을 취소했다.

 

isinstance는 데이터타입이나 인스턴스를 확인하는 함수로 isinstance(변수, 타입) 형태로 쓴다.

ex)

x=10, y=4.5

isinstance(x, int) -> x는 정수인가?  true

isinstance(y,(int, float) -> y는 (정수, 실수) 다음 튜플 중에 해당하는가? -> 실수이므로 true

 

class Dog:
   pass

 

Bob=Dog()

isinstance(Bob, Dog) -> Bob은 Dog 클래스의 인스턴스인가? -> true

 

 if isinstance(layer, nn.ReLU) -> layer은 nn.ReLu 클래스의 인스턴스인가? 

이 값이 True일때만 덮어쓰기가 되지 않게 했다.

 

여기서 잠깐 vgg16의 코드예시와 구조를 보고 가면

class VGG16(nn.Module):
    def __init__(self, num_classes=n):  
        super(VGG16, self).__init__()
        
        # Convolutional Layers
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
                         중략
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # Fully Connected Layers
        self.classifier = nn.Sequential(
            nn.Linear(512 * 1 * 1, 4096),
            nn.ReLU(True),
                         중략
            nn.Linear(4096, num_classes)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # Flatten
        x = self.classifier(x)
        return x

이렇게 conv, relu, maxpool이 반복되며 깊은 층을 형성하고 있다.

 

물론 vgg16을 쓸때는 저 긴 정의코드를 다 쓰지 않고 단순히 모델을 불러오기만 하면 된다.

vgg16 = models.vgg16(pretrained=True).features

바로 이렇게

코드에서 보면 알 수 있듯이 모델의 FC layer나 순전파는 사용하지 않고 features 만 가져오고 있다.

style transfer는 이미지를 분류하는 작업이 아니기 때문에 분류부분은 굳이 필요가 없다.

 

나중에도 features 라는 말이 나올텐데, 바로 저 vgg16.features에 해당하는 convolution layer 부분과 관계가 깊다.

 

def gram_matrix(features):
    with torch.no_grad():
        _, C, H, W = features.size()
        features = features.view(C, H * W)
        gram = torch.mm(features, features.t()) / (C * H * W)
    return gram

gram_matrix를 정의하는 부분인데, (수학적인 내용은 간단하게만)

gram행렬 내의 요소는 채널간의 유사도를 반영하므로, 두 gram matrix 값이 높으면 둘의 상관관계가 크다는 의미

특징맵(features)에서 채널수, 높이, 너비를 각각 뽑아내고 2D텐서로 변환 후 gram matrix를 계산하고 있다.

 

gram = torch.mm(features, features.t()) / (C * H * W)

이 부분이 바로 gram matrix를 계산해내는 부분

2D텐서로 펴진 features와 그의 전치행렬인 features.t()간 행렬곱을 수행하고, C * H * W 로 나누어 정규화 해주고 있다.

정규화를 함으로써 이미지의 크기나 채널수와 관계없이 gram matrix를 계산할 수 있다.

 

참고로 torch.no_grad도 이후 역전파 부분에서 자꾸 오류가 나서 추가했다.

gram matrix는 스타일 손실 계산시에만 필요하기 때문에 역전파에서는 무시할 수 있다.

 

def get_features(image, model, layer_idx):
    content_feature = None
    style_features = []
    x = image
    
    for i, layer in enumerate(model):
        x = layer(x)
        
        if i == layer_idx['content']:
            content_feature = x.clone()
        elif i in layer_idx['style']:
            style_features.append(gram_matrix(x))
    
    return content_feature, style_features

 

layer_idx = {'content': 21, 'style': [0, 5, 7, 10, 16, 24, 26]}

 

with torch.no_grad():
    content_features_from_content, style_features_from_content = get_features(content_image, vgg16, layer_idx)
    content_features_from_style, style_features_from_style = get_features(style_image, vgg16, layer_idx)

 

특징맵 추출함수를 정의하고 함수로 content image와 style image에서 특징맵 추출

 

layer_idx = {'content': 21, 'style': [0, 5, 7, 10, 16, 24, 26]}

이 부분을 먼저 보겠다.

콘텐츠 레이어와 스타일 레이어를 지정해주는 코드인데, 위의 vgg모델 features 부분을 참고해보면 

21번 레이어는 다음과 같은 conv 층이다.

21:Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

 

0, 5, 7, 10, 16, 24, 26도 각각 conv층이다. 

vgg16.features[idx] 를 입력해도 각 층을 확인할 수 있다. 

주어진 이미지에서 contents feature를 뽑아낼 층과 style feature를 뽑아낼 층들을 지정해주고 있다.

이 값은 사용하는 vgg 모델에 따라 다를 수 있으니 직접 보고 정해야 한다.

content layer는 이미지의 구도를 유지하는 것이 목적이므로 구도를 담고있는 깊은층 하나를 지정하고,

style layer는 얕은층부터 깊은층까지 다양하게 지정하여 여러 특징들을 특징맵으로 만들고 있다.

 

    content_feature = None
    style_features = []
content feature는 하나의 conv 층에서 나온 특징맵이므로 단일값이고, style feature는 여러개의 layer에서 추출한 특징맵이므로 리스트형태이다. 여기에 이후 뽑아낸 특징맵이 저장된다.

 

    for i, layer in enumerate(model):
        x = layer(x)

 

이미지x가 지정된 레이어를 통과하면서 값을 반환한다. 

x=layer(x) 란 이미지가 예를 들어 convolutional layer 로 들어가면 합성곱을 수행하고 특징맵을 만든다는 뜻

반복문이므로 이렇게 만든 특징맵을 x에 저장하고, 다음 레이어로 넘겨주게 된다.

 

        if i == layer_idx['content']:
            content_feature = x.clone()

컨텐츠레이어에 해당하는 인덱스라면, x의 clone을 만들고 그것을 content feature에 저장함

원래코드는 content_feature=x 였는데 나중에 오류가 많이 나서 독립적으로 clone을 만들어 안정성을 높였다. 

 

        elif i in layer_idx['style']:
            style_features.append(gram_matrix(x))

 

style idx는 리스트 형태이기 때문에 if i==[리스트] 로 쓰면 이 조건문은 항상 false가 된다. 그래서 i in [리스트]

그리고 style_features=[] 인 빈 리스트에 gram matrix를 저장하고 있다.

참고로 elif 로 썼으므로 content layer와  style layer의 idx번호가 같다면 content feature에만 저장된다.

두 layer에서 같은 idx를 사용하면 content feature와 style feature가 유사해질 수 있다. 더 좋을지 나쁠지는..?

 

with torch.no_grad():
    content_features_from_content, style_features_from_content = get_features(content_image, vgg16, layer_idx)
    content_features_from_style, style_features_from_style = get_features(style_image, vgg16, layer_idx)

 

이제 정의한 get_features함수를 사용해 content image와 style image에서 content feature와 style feature를 생성했다.

각각의 feature는 손실함수를 계산하는데 사용된다.

 

generated_image = content_image.clone().requires_grad_(True)

초기 생성이미지

 

이 모델은 content image vs generated image -> 손실을 계산하고 가중치 업그레이드 -> content image vs generated image with updated weight -> 둘의 손실을 또 계산하고 가중치 업그레이드 ->... 손실을 줄이는 방향으로 나아감

콘텐츠 손실이 줄어든다 = content image와 generated image간 구도와 모양이 유사해진다.

스타일 손실이 줄어든다= style image와 generated image간 질감, 색상, 엣지, 형태가 모두 유사해진다.

 

이 때 가장 처음에는 generated image가 없으니 content image를 복사해서 임시로 generated image로 사용

 

Q. 그럼 거의 동일한 두 이미지를 비교해서 손실을 계산한다는건데 손실이 0 아닌가?

A. 스타일 손실은 style image와 generated image간의 차이이니 0이 아니다. 또한 콘텐츠 손실 역시 두 이미지가 합성곱 레이어를 통과하면서 미세한 차이가 발생하여 0은 아니게 된다.

뒤에서 손실에 대해 더 자세히 이야기하는걸로

 

requires_gard_(True) 도 역전파 과정에서 오류가 나서 추가한 코드인데, 텐서에 대한 손실함수의 그라디언트가 계산되어 역전파 과정에서 최적화가 진행된다고 한다.(솔직히 이건 무슨 말인지 잘 모르겠다.)

 

content_weight = 10
style_weight = 5000

# 학습 루프
num_iterations = 1000
for i in range(num_iterations):
    optimizer.zero_grad()
    
    content_features_from_generated, style_features_from_generated = get_features(generated_image, vgg16, layer_idx)
    

    content_loss = criterion(content_features_from_content, content_features_from_generated)
   
    style_loss = sum(criterion(style_gram, generated_gram)
                     for style_gram, generated_gram in zip(style_features_from_style, style_features_from_generated))
    
    total_loss = content_weight * content_loss + style_weight * style_loss
    
    print(f"Iteration {i+1}/{num_iterations}, Total Loss: {total_loss.item()}")
    
    total_loss.backward()
    optimizer.step()

    with torch.no_grad():
        generated_image.clamp_(0, 1)

content_weight = 10  #0.1, 1, 10 등 작은값
style_weight = 5000  #1e3, 1e4 등 큰값

이 weight은 total_loss = content_weight * content_loss + style_weight * style_loss 에서 사용하는 가중치인데, 

모델을 돌려보면서 최적의 값으로 본인이 설정한다.

content weight이 클수록 content image의 특성이 유지되고, style weight이 클수록 style image의 특성을 더 많이 반영하므로 보통 style weight이 훨씬 크다. 

 

 content_loss = criterion(content_features_from_content, content_features_from_generated)

위에서도 이야기했듯 content 손실은 단순히 두 이미지간의 content feature를 비교하면 된다. 여긴 간단하다.

 

    style_loss = sum(criterion(style_gram, generated_gram)
                     for style_gram, generated_gram in zip(style_features_from_style, style_features_from_generated))

스타일손실이 복잡한데, 스타일손실은 스타일레이어의 특징맵으로 구한 그람메트릭스끼리의 손실의 합

ex)

A=스타일그림 , B=생성된그림

style layer 번호= [5,10,16,24]

위로 올려보면 style features란 레이어에서 뽑아낸 Gram Matrix가 들어있는 리스트이다.

A에서 뽑아낸 style features_A= [GMA1, GMA2, GMA3, GMA4]

B에서 뽑아낸 style features_B= [GMB1, GMB2, GMB3, GMB4]

zip은 두 리스트를 병렬로 튜플로 묶어주는 함수

zip( style features_A , style features_B) =[( GMA1, GMB1 ),( GMA2 , GMB2 ),( GMA3 , GMB3 ),( GMA4 , GMB4 )]

이들의 손실함수를 각각 구하면

criterion ( GMA1, GMB1 ), criterion ( GMA2, GMB2 ), criterion ( GMA3, GMB3 ), criterion ( GMA4, GMB4 )

이들을 for in 문으로 반복했고 합해주는 함수 sum을  썼다.

 

     total_loss = content_weight * content_loss + style_weight * style_loss

총 손실은 두 손실을 가중치 곱해서 더하면 된다.

 

    with torch.no_grad():
        generated_image.clamp_(0, 1)

.clamp_(0,1) 

clamp_는 텐서의 값이 지정된 범위 안에 있게 하고, in-place 연산으로 메모리를 추가로 쓰지 않고 텐서값을 직접 수정한다. 이 때 지정범위는 0에서 1 사이이다. 딥러닝 모델을 대게 이미지 픽셀을 이 사이로 정규화하기 때문

이미지가 모델을 거치면서 여러 이유로 픽셀값이 음수나 1보다 큰 수가 되기도 하는데, 시각화할 때 오류를 피하기 위해 픽셀을 모두 0과 1 사이에 들어오도록 맞춰주는것

이것도 원래 없었는데 오류나서 추가했다...

 

    if (i + 1) % 5 == 0 or i == num_iterations - 1:
        plt.figure()
        unloader = transforms.ToPILImage()
        image = generated_image.cpu().clone().squeeze(0)  
        image = unloader(image)
        plt.imshow(image)
        plt.title(f"Generated Image at Iteration {i+1}")
        plt.axis('off')
        plt.show()

이미지를 생성해주는 코드인데

i == num_iterations - 1 를 추가해서 마지막 실행헤서 최종결과물을 내도록 했고

image = generated_image.cpu().clone().squeeze(0)  이미지 생성을 위해 cpu로 가져오는 코드도 넣었다.

 

결과?!

content image

chatgpt한테 요청함

 

style image

이건 chatgpt한테 요청한것이고 실제 이미지는 이것과 유사한 느낌의 아크릴화 사용

 

 

알록달록한 강아지를 기대했는데.. 그래도 노력은 했다

loss 거의 변화없어 중지시킴

 

문제점?

1. 위는 수정된 코드인데, style layer를 처음에 조금 잘못 설정해서 다양한 특징을 학습하지 못했을 수 있다. 

2. 최적화알고리즘 : 현재 Adam  사용했는데, LBFGS가 더 자주 사용되고 성능도 좋다고 한다.

 

다음번에 LBFGS를 사용한 결과를 들고 오도록 하겠다..

참고로 LBFGS는 gpu에서 동작이 아주 느리니 차라리 cpu에서 쓰도록  하자