스타일, 콘텐츠 및 변형 손실 계산하기

스타일, 콘텐츠 및 변형 손실 계산하기

 

 

행렬에 전치를 곱하여 그램 행렬을 계산한다. 손실 함수의 두 부분을 계산하기 위해 VGG 네트워크의 여러 컨볼루션 레이어에서 출력된 그램 행렬을 가져온다. 스타일과 원본 이미지와의 유사성을 모두 결정하기 위해 이미지 픽셀을 직접 비교하는 대신 VGG의 컨볼루션 레이어 출력을 비교한다. 손실 함수의 세 번째 부분에서는 서로 가까운 픽셀을 직접 비교한다.

 

VGG 네트워크의 여러 다른 레벨에서 컨볼루션 출력을 가져오기 때문에 그램 행렬은 이러한 레이어를 결합하는 수단을 제공한다. VGG 컨볼루션 레이어의 그램 행렬은 이미지의 스타일을 나타낸다. 알고리즘이 생성할 때 원본 이미지, 스타일 참조 이미지, 최종 출력 이미지에 대해 이 스타일을 계산한다.

 

def gram_matrix(x):
    x = tf.transpose(x, (2, 0, 1))
    features = tf.reshape(x, (tf.shape(x)[0], -1))
    gram = tf.matmul(features, tf.transpose(features))

    return gram

def style_loss(style, combination):
    S = gram_matrix(style)
    C = gram_matrix(combination)
    channels = 3
    size = img_nrows * img_ncols

    return tf.reduce_sum(tf.square(S - C)) / (4.0 * (channels ** 2) * (size ** 2))

def content_loss(base, combination):
    return tf.reduce_sum(tf.square(combination - base))

def total_variation_loss(x):
    a = tf.square(
        x[:, :img_nrows - 1, :img_ncols - 1, :] \
        - x[:, 1:, :img_ncols - 1, :]
    )
    b = tf.square(
        x[:, :img_nrows - 1, :img_ncols - 1, :] \
        - x[:, :img_nrows - 1, 1:, :]
    )
    
    return tf.reduce_sum(tf.pow(a + b, 1.25))

 

style_loss 함수는 현재 생성된 이미지 (조합)가 참조 스타일 이미지의 스타일과 얼마나 일치하는지 비교한다. 스타일과 현재 생성된 이미지의 그램 행렬을 차감하고 정규화하여 이 스타일의 차이를 계산한다. 정확하게는 서로 다른 VGG 레이어에서 추출한 기본 이미지와 스타일 참조 이미지의 표현의 그램 행렬 사이의 L2 거리의 합으로 구성한다. 일반적인 아이디어는 다양한 공간 스케일 (고려되는 레이어의 깊이에 따라 정의되는 상당히 큰 스케일)에서 색상/텍스처 정보를 캡처하는 것이다.

 

content_loss 함수는 현재 생성된 이미지가 원본 이미지와 얼마나 일치하는지 비교한다. 이 차이를 계산하려면 원본 이미지와 생성된 이미지의 그램 행렬을 빼야 한다. 여기서는 기본 이미지의 VGG 특징과 생성된 이미지의 특징 사이의 L2 거리를 계산하여 생성된 이미지가 원본 이미지에 충분히 가깝게 유지되도록 한다.

 

마지막으로, total_variation_loss 함수는 생성된 이미지의 픽셀 사이에 국소적인 공간 연속성을 부여하여 시각적 일관성을 부여한다.

 

The VGG Neural Network

 

VGG19는 K. Simonyan과 A. Zisserman이 제안한 컨볼루션 신경망 모델이다. 이 모델은 1000개 클래스에 속하는 1400만 개 이상의 이미지로 구성된 데이터 세트인 ImageNet에서 92.7%의 상위 5위 테스트 정확도를 달성했다. VGG16 가중치를 스타일 전송 모델로 옮긴다. Keras는 VGG 신경망을 로드하는 기능을 제공한다.

 

model = vgg19.VGG19(weights="imagenet", include_top=False)
outputs_dict = dict([(layer.name, layer.output) for layer in model.layers])
feature_extractor = keras.Model(inputs=model.inputs, outputs=outputs_dict)

 

이제, 완전한 손실 함수를 생성할 수 있다. 다음 이미지가 compute_loss 함수에 입력된다.

 

combination_image : 생성된 이미지의 현재 반복 횟수 base_image : 시작 이미지 style_reference_image : 재현할 스타일이 담긴 이미지

 

style_layer_names으로 지정된 레이어는 세 이미지 각각에 대해 VGG에서 피처로 추출할 레이어를 나타낸다.

 

style_layer_names = [
    "block1_conv1",
    "block2_conv1",
    "block3_conv1",
    "block4_conv1",
    "block5_conv1",
]

content_layer_name = "block5_conv2"

def compute_loss(combination_image, base_image, style_reference_image):
    input_tensor = tf.concat(
        [base_image, style_reference_image, combination_image], axis=0
    )
    features = feature_extractor(input_tensor)
    
    loss = tf.zeros(shape=())
    
    layer_features = features[content_layer_name]
    base_image_features = layer_features[0, :, :, :]
    combination_features = layer_features[2, :, :, :]
    
    loss = loss + content_weight * content_loss(
        base_image_features, combination_features
    )
    
    for layer_name in style_layer_names:
        layer_features = features[layer_name]
        style_reference_features = layer_features[1, :, :, :]
        combination_features = layer_features[2, :, :, :]
        sl = style_loss(style_reference_features, combination_features)
        loss += (style_weight / len(style_layer_names)) * sl
    
    loss += total_variation_weight * total_variation_loss(combination_image)
    
    return loss

 

스타일 전송 이미지 생성

 

compute_loss_and_grads 함수는 손실 함수를 호출하고 그라데이션을 계산한다. 이 모델의 매개변수는 생성된 이미지의 현재 반복에 대한 실제 RGB 값이다. 이러한 매개변수는 기본 이미지에서 시작하여 알고리즘이 최종 렌더링된 이미지에 맞게 최적화한다. 변환을 수행하기 위해 모델을 훈련시키는 것이 아니라 손실 함수를 최소화하기 위해 이미지를 훈련/수정한다. 신경망 훈련이 가중치를 수정하는 것과 같은 방식으로 Keras가 이미지를 수정할 수 있도록 gradient tape를 활용한다.

 

def compute_loss_and_grads(combination_image, base_image, style_reference_image):
    with tf.GradientTape() as tape:
        loss = compute_loss(combination_image, base_image, style_reference_image)
        grads = tape.gradient(loss, combination_image)
        
    return loss, grads

 

이제, 손실 함수에 따라 이미지를 최적화할 수 있다.

 

optimizer = keras.optimizers.SGD(
    keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate=100.0, decay_steps=100, decay_rate=0.96
    )
)

base_image = preprocess_image(base_image_path)
style_reference_image = preprocess_image(style_reference_image_path)
combination_image = tf.Variable(preprocess_image(base_image_path))
iterations = 5000

for i in range(1, iterations + 1):
    loss, grads = compute_loss_and_grads(
        combination_image, base_image, style_reference_image
    )

    optimizer.apply_gradients([(grads, combination_image)])
    if i % 100 == 0:
        print("Iteration %d: loss=%.2f" % (i, loss))
    img = deprocess_image(combination_image.numpy())
    fname = result_prefix + "_at_iteration_%d.png" % i
    keras.preprocessing.image.save_img(fname, img)

display(Image(result_prefix + "_at_iteration_5000.png"))

from google.colab import files

files.download(result_prefix + "_at_iteration_5000.png")