[Paper Explain] DualStyleGAN: Exemplar-Based High-Resolution Portrait Style Transfer

    image.png

Introduction

    Gần đây, nhiều nghiên cứu chỉ ra rằng StyleGAN có thể thực hiện style transfer chất lượng cao chỉ với một lượng dữ liệu hạn chế bằng một chiến lược fine tuning phù hợp. Paper Pastiche Master: Exemplar-Based High-Resolution Portrait Style Transfer đề xuất một mở rộng của kiến trúc StyleGAN với intrinsic style path và extrinsic style path để mã hóa style của domain gốc và domain cần transfer sang. Paper cũng đề xuất một chiến lược fine tuning theo 3 bước để biến đổi không gian của domain gốc sang domain mới ngay cả với những thay đổi về kiến trúc mạng.

Một số khái niệm

StyleGAN generator

    Mạng generator của stylegan có 1 chút khác biệt với các mô hình GAN truyền thống. Thay vì nhận đầu vào trực tiếp là vector latent ziZz_i \in Z thì vector zz sẽ được đưa qua một mạng MLP 8 lớp để tạo ra vector wWw \in W với cùng số chiều. Vector ww được dùng để kiểm soát style của ảnh thông qua lớp Adaptive Instance Normalization. Một phép biến đối affline (học thông qua một lớp fully connected - ký hiệu A trong hình dưới) được áp dụng lên ww trước khi đưa vào mạng generator.

    

GAN inversion

    Mạng generator của GAN sẽ tạo ra ảnh từ một vector latent zz. GAN inversion là quá trình ngược lại, từ ảnh đầu vào tìm vector latent tương ứng của nó.

    image.png

    Một số cách tiếp cận cho GAN inversion

  • Optimization based: sử dụng gradient descent để tối ưu vector zz sao cho ảnh khôi phục xrecx^{rec} giống với ảnh thật xx nhất.

    z=argminz(x,G(z;θ))z ^ { * } = \underset { z } { \arg \min } \ell ( x , G ( z ; \theta ) )

  • Encoder based: sử dụng 1 mạng encoder được huấn luyện trên nhiều ảnh được tạo ra bởi mạng generator với các latent vector tương ứng.

    θE=argminθEnL(G(E(xn;θE)),xn)\theta _ { E } ^ { * } = \underset { \theta _ { E } } { \arg \min } \sum _ { n } L ( G ( E ( x _ { n } ; \theta _ { E } ) ) , x _ { n } )

    Trong paper này tác giả sử dụng encoder based dựa trên model psp

Style mixing

    Trong lúc huấn luyện StyleGAN, một phần ảnh trong tập training được tạo ra bằng cách sample 2 latent vector z1z_1z2z_2 thay vì một. Trong quá trình forward của mạng generator, ta lựa chọn ngẫu nhiên một điểm trong mạng để chuyển từ vector style w1w_1 sang w2w_2.

Dual StyleGAN

    DualStyleGAN được xây dựng dựa trên một mạng StyleGAN pretrain. Việc finetune lại mô hình unconditional trên một dataset khác sẽ làm dịch chuyển cả không gian sample, dẫn tới mất đa dạng trong các sample sinh ra. Ý tưởng chính của DualStyleGAN tìm kiếm chiến lược finetuning phù hợp để học các style đa dạng bằng cách thêm extrinsic style path để encode style của domain mục tiêu image.png

Facial Destylization

    Mục tiêu của bước này là khôi phục ảnh người thật từ ảnh style để tạo pair dataset huấn luyện mô hình. Với một ảnh style từ domain mục tiêu, ta cần tìm một khuôn mặt tương ứng phù hợp ở domain gốc. Quá trình này được thực hiện qua 3 bước:

  1. Khởi tạo: Một ảnh style S đầu tiên sẽ được embed thành 1 vector latent trong domain gốc thông qua quá trình GAN inversion. Cụ thể ở đây tác giả sử dụng mô hình psp kí hiệu là EE. Thay vì embed ảnh thành vector trong không gian Z\mathcal Z (1 vector sample từ phân phối chuẩn), tác giả sử dụng không gian Z+\mathcal Z^+ (18 vector sample từ phân phối chuẩn, stack lên thành ma trận 18×51218 \times 512).

    Khuôn mặt sau khi tái tạo trong domain gốc kí hiệu là g(ze+)g(\pmb z_e^+). Với gg là generator của mạng StyleGAN gốc và ze+=E(S)R18×512\pmb z_e^+ = E(S) \in \mathbb R^{18 \times 512}.

  1. Finetune vector latent: Không gian Z+\mathcal Z^+ của một mạng StyleGAN đẫ được finetune trên tập ảnh style (ký hiệu gg') được tối ưu để sinh ra ảnh style bằng cách finetune StyleGAN với hàm loss

    z^e+=argminz+Lperc(g(z+),S)+λIDLID(g(z+),S)+σ(z+)1\hat { \mathbf { z } } _ { e } ^ { + } = \arg \min _ { z ^ { + } } L _ { { p e r c } } \left ( g ^ { \prime } \left ( z ^ { + } \right ), S \right ) + \lambda _ { ID } L _ { ID } \left ( g ^ { \prime } \left ( z ^ { + } \right ), S \right ) + \| \sigma \left ( z ^ { + } \right ) \| _ { 1 }

    trong đó

  • LpercL_{perc}: là perceptual loss với đặc trưng trích xuất từ VGG19
  • LID=1cos(F(g(z+)),F(S))L_{ID} = 1 - \cos(F( g ^ { \prime } \left ( z ^ { + } \right )), F(S)): identity loss để bảo toán nhận dạng với FF là mô hình Arcface.
  • σ(z+)1\|\sigma \left (z^{ + } \right ) \| _ { 1 }: regularization term - độ lệch chuẩn giữa 18 vector trong z+z^+. Dùng để giảm overfitting
  1. Image embedding : Vector embedding cuối cùng được cho bằng công thức zi+=E(g(z^e+))z _ { i } ^ { + } = E \left ( g \left ( \hat { z } _ { e } ^ { + } \right ) \right ). Kết quả reconstruct qua 3 bước finetune được cho trong hình (d)

    image.png

Kiến trúc mạng

    image.png

    Kiến trúc mạng DualStyleGAN bao gồm 2 phần gồm 2 phần:

  • intrinsic style path + generator. Phần này là kiến trúc StyleGAN gốc và được fix cứng trong quá trình finetune.
  • Extrinsic style path: gồm 1 mạng MLP giống intrinsic path, các khối color transform TcT_c và structure transform TsT_s. Đầu vào của extrinsic path là vector ze+\boldsymbol z_e^+ của ảnh style tương ứng.

    Color control block: phần này khá giống StyleGAN gốc. Style code từ extrinsic path đi qua một mạng MLP ff và các khối color transform TcT_c được mô hình hóa bởi 1 mạng MLP. Các khối này chỉ nằm ở phần có độ phân giải cao ở generator (lớp 8-18). Style code sau khi đi qua TcT_c sẽ được cộng có trọng số với style code từ intrinsic path trước khi đi qua Adaptive Normalization.

    Structure control block: dùng để đặc trưng hóa cấu trúc của ảnh style. Style code sau khi đi qua mạng ff và structure control block TsT_s sẽ được cho qua 1 khối ModRes trước khi được cộng với feature map của generator. Các khối này chỉ xuất hiện ở các lớp độ phân giải thấp (lớp 1 - 7). Các trúc của khối ModRes giống hệt các residual block trong Resnet chỉ khác là các lớp Batch normalization được thay thế bằng Adaptive Instance Normalization.

    Với một ảnh mặt thật II và ảnh style SS, style transfer được thực hiện bằng G(E(I),E(S),w)G(E(I), E(S), w) với wR18w \in \mathbb R^{18} và trọng số để kết hợp giữa style của intrinsic và extrinsic path.

Progressive finetuning

    DualStyleGAN sử dụng chiến lược finetuning nhiều bước để từ từ biến đổi không gian sample sang domain mục tiêu. Hai bước đầu có mục đích là để pretrain DualStyleGAN trên data gốc sau đó mới finetune trên data mới trong bước 3. image.png

Bước 1: color transfer

    Nhờ vào thiết kế của DualStyleGAN, bước này ta không cần huấn luyện lại mạng mà chỉ cần khởi tạo các khối ở extrinsic path một cách hợp lý.

  • Các khối ModRes được khởi tạo gần 0 để loại bỏ các residual feature
  • Trọng số các khối color transform được khởi tạo bằng ma trận đơn vị nên style code không thay đổi khi đi vào generator. Có thể thấy màu của ảnh style được transfer khá tốt sang ảnh mặt người chỉ nhờ khởi tạo trong hình trên

Bước 2: structure transfer

    Finetune extrinsic path để transfer các đặc trưng về structure

    Đầu tiên ta sample 2 vector z1,z2z_1, z_2 từ phần phối chuẩn. DualStyleGAN được finetune sao cho G(z1,z2,w=1)g(zl+)G(z_1, z_2, w=1) \approx g(z_l^+). Trong đó, zl+Z+z_l^+ \in \mathcal Z^+ là ma trận 18×51218 \times 512 với ll hàng đầu tiên là z1z_118l18-l hàng cuối là z2z_2. ll giàm dần từ 7 xuống 5 trong quá trình finetune. Hàm mục tiêu:

    minGmaxDλadvLadv+λpercLperc(G(z1,z2,w=1),g(zl+))\min _ { G } \max _ { D } \lambda _ { \operatorname { a d v } }L_{adv} + \lambda_{perc}L_{perc}(G(z_1, z_2, w=1), g(z_l^+))

Bước 3: finetune trên domain mục tiêu

    Ta cần style code zi+z_i^+ze+z_e^+ của một ảnh style S trong domain mục tiêu có thể được tái tạo thông qua DualStyleGAN G(zi+,ze+,1)SG(z_i^+, z_e^+, 1) \approx S thông qua perceptual loss. Hàm mục tiêu của bước 3:

    minGmaxDλadvLadv+λpercLperc(G(zi+,ze+,1),S)+Lsty+Lcon\min _ { G } \max _ { D } \lambda _ { \operatorname { a d v } }L_{adv} + \lambda_{perc}L_{perc}(G(z_i^+, z_e^+, 1), S) + L_{sty} + L_{con}

    trong đó

    Lsty=λCXLCX(G(z,ze+,1),S)+λFMLFM(G(z,ze+,1),S)L_{sty} = \lambda_{CX}L_{CX}(G(z, z_e^+, 1), S) + \lambda_{FM}L_{FM}(G(z, z_e^+, 1), S)

    LCXL_{CX}LFML_{FM} lần lượt là context lossfeature matching loss

    Lcon=λIDLID(G(z,ze+,1),g(z))+λregW2L_{con} = \lambda_{ID}L_{ID}(G(z, z_e^+, 1), g(z)) + \lambda_{reg} \| W\|_2

    với WW là trọng số của các lớp ModRes.

Result

    Một số kết quả thí nghiệm image.png

    image.png

Reference

Nguồn: Viblo

Bình luận
Vui lòng đăng nhập để bình luận
Một số bài viết liên quan