[Unity Barracuda] 유니티 바라쿠다 튜토리얼 (StyleTransfer-AdaIN)

 

---

Barracuda

Unity Technologies

---

 

서론

 

많이 늦어졌지만 Unity Barracuda를 실습해보고자 합니다. 기존 사내에 하늘 같은 그 분께서 AdaIN을 사용하시는 것을 보여주셨는데 해당 내용을 토대로 내용을 진행해보겠습니다.

 

본문

 

모델 선정 naoto0804/Pytorch-AdaIN 

GitHub - naoto0804/pytorch-AdaIN: Unofficial pytorch implementation of 'Arbitrary Style Transfer in Real-time with Adaptive Inst

 

선정 기준은 따로 존재하지 않습니다. Github에서 파이썬으로 가장 많이 star를 받은 repo를 기준으로 선정하였습니다.

 

 

모델 확인 (Python)

 

실제로 정상적으로 해당 모델이 작동하는지 확인해보겠습니다. 사용 방법은 Github에 정리가 되어 있어서 쉽게 따라할 수 있었습니다. 현재 블로그에 있는 프로필 사진을 토대로 만들어보겠습니다.

 

python test.py --content input/content/profile.jpg --style input/style/antimonocromatismo.jpg

 

그 결과는 아래와 같이 출력되었습니다. 마음에 들지는 않지만... 변환이 잘 되는 것이 중요하니 우선 진행해보도록 하겠습니다. 참고로 저는 사진의 크기를 512x512로 맞추어서 진행하였습니다.

 

 

모델의 Shape 확인하기

 

Onnx로 export 하기위해서는 모델의 shape을 먼저 확인해야 합니다. 이 경우 512x512로 맞춰서 진행하였기 때문에 큰 문제가 없지만 코드에서 확인하면 다음과 같습니다.

모델의 Shape는 추론모드에서 확인을 하기 때문에  eval(), device 할당 및 unsqueeze 다음에 확인이 가능하며 코드 하단 부에 one content image, N style image 그리고 one content and one style 이렇게 두 부분으로 나누어져 있는데 저는 one content and one style로 진행하였기 때문에 else 문에서 확인하였습니다.

 

for content_path in content_paths:
    if do_interpolation:  # one content image, N style image
        style = torch.stack([style_tf(Image.open(str(p))) for p in style_paths])
        content = content_tf(Image.open(str(content_path))) \
            .unsqueeze(0).expand_as(style)
        style = style.to(device)
        content = content.to(device)

        with torch.no_grad():
            output = style_transfer(vgg, decoder, content, style,
                                    args.alpha, interpolation_weights)
        output = output.cpu()
        output_name = output_dir / '{:s}_interpolation{:s}'.format(
            content_path.stem, args.save_ext)
        save_image(output, str(output_name))

    else:  # process one content and one style
        for style_path in style_paths:
            content = content_tf(Image.open(str(content_path)))
            style = style_tf(Image.open(str(style_path)))
            if args.preserve_color:
                style = coral(style, content)
            style = style.to(device).unsqueeze(0)
            content = content.to(device).unsqueeze(0)
            
            # Print Shapes!
            print("Content shape:", content.shape)
            print("Style shape:", style.shape)
            
            with torch.no_grad():
                output = style_transfer(vgg, decoder, content, style,
                                        args.alpha)
            output = output.cpu()

            output_name = output_dir / '{:s}_stylized_{:s}{:s}'.format(
                content_path.stem, style_path.stem, args.save_ext)
            save_image(output, str(output_name))

 

중간에 #Print Shapes! 이 후 2 줄을 보시면 확인이 가능합니다. 그 결과 예상하던 것과 같이다음과 같이 확인이 가능합니다.

 

 

Onnx Export

 

import torch와 import torchvision은 이미 되어 있음으로 별도로 수정하지 않았습니다.

Onnx로 export 하는 내용은 파이토치 사이트에 나와 있어서 해당 내용을 참고하였습니다. (본문 링크)

(선택) PyTorch 모델을 ONNX으로 변환하고 ONNX 런타임에서 실행하기

 

바로 밑에 있는 style_transfer에 있는 코드 내용을 그대로 사용하면 됩니다. 우리는 interpolation을 어떻게 사용할지 이미 알고 있으니 해당 슬라이싱 부분은 제외하고 진행하였습니다.

 

class AdaIN(nn.Module):
    def __init__(self, vgg, decoder, alpha):
        super(AdaIN, self).__init__()
        self.vgg = vgg
        self.decoder = decoder
        self.alpha = alpha

    def forward(self, content, style, alpha=1.0, interpolation_weights=None):
        return style_transfer(self.vgg, self.decoder, content, style, self.alpha, interpolation_weights)

 

클래스 구조체를 만들어줍니다. 이 후 하단 부에 AdaIN을 호출하여 adain에 넣어줍니다. 이 후 기기설정 및 평가모드로 전환하고 onnx로 export 해줍니다. 추가한 내용은 다음과 같습니다.

 

adain = AdaIN(vgg, decoder)
            adain.to(device)
            adain.eval()
            torch.onnx.export(adain, (content, style), "adain.onnx", opset_version=9, input_names=['content', 'style'], output_names=['output'])

 

전체 코드 본문에서는 다음과 같습니다.

 

    else:  # process one content and one style
        for style_path in style_paths:
            content = content_tf(Image.open(str(content_path)))
            style = style_tf(Image.open(str(style_path)))
            if args.preserve_color:
                style = coral(style, content)
            style = style.to(device).unsqueeze(0)
            content = content.to(device).unsqueeze(0)
            
            adain = AdaIN(vgg, decoder)
            adain.to(device)
            adain.eval()
            torch.onnx.export(adain, (content, style), "adain.onnx", opset_version=9, input_names=['content', 'style'], output_names=['output'])

            # Print the shapes
            print("Content shape:", content.shape)
            print("Style shape:", style.shape)
            
            with torch.no_grad():
                output = style_transfer(vgg, decoder, content, style,
                                        args.alpha)
            output = output.cpu()
            
           

            output_name = output_dir / '{:s}_stylized_{:s}{:s}'.format(
                content_path.stem, style_path.stem, args.save_ext)
            save_image(output, str(output_name))

 

Onnx Inference

 

그 분께서 강력하게 권장하셨던 내용입니다. 우선 모델이 정상 작동하는지 확인하고, 모델 결과를 저장하고, onnx로 내보내고 마지막으로 onnx inference를 꼭! 해봐야 추후에 결과가 이상하게 나와도 뭐가 잘못됐는지 알 수 있다고 말씀하셨습니다. 사실 이렇게 step별로 진행하는게 매우 중요하다는걸... 항상 느끼고 있습니다.

 

이 부분에서 생각보다 많이 해맸었네요. 관련한 내용은 파이토치 사이트에 정리되어 있습니다.

(선택) PyTorch 모델을 ONNX으로 변환하고 ONNX 런타임에서 실행하기

 

주된 내용 PIL 라이브러리를 이용하는데 이 때 export 된 onnx는 tensor 값만을 입력으로 받기 때문에 이미지를 텐서로 변환해줘야 합니다. 전체 코드는 다음과 같습니다.

 

import onnx
import onnxruntime
import torch
from PIL import Image
import numpy as np
from torchvision import transforms

#모델 불러오기
AdaIN = onnxruntime.InferenceSession("adain.onnx") 

#각 이미지 경로 불러오기
content_img = Image.open("input/content/profile.jpg")
style_img = Image.open("input/style/antimonocromatismo.jpg")

#텐서 변환을 위한 to_tensor
to_tensor = transforms.ToTensor()

#배치 차원 추가
content_img_tensor = to_tensor(content_img).unsqueeze(0)
style_img_tensor = to_tensor(style_img).unsqueeze(0)

#Adain.onnx에 넣을 Input 셋팅
content_inputs = {AdaIN.get_inputs()[0].name: content_img_tensor.numpy()}
style_inputs = {AdaIN.get_inputs()[1].name: style_img_tensor.numpy()}

#onnx inference
output_data = AdaIN.run(None, {**content_inputs, **style_inputs})

#결과 값을 tensor로 변환
output_tensor = torch.from_numpy(output_data[0])
output_tensor = output_tensor.squeeze(0)

#이미지로 저장
output_img = transforms.ToPILImage()(output_tensor.clamp(0, 1))
output_img.save("output/result__.jpg")

 

이렇게 한 결과는 다음과 같습니다. Offset 버전이 달라서 걱정 하였지만 큰 차이 없이 변변환이 되었습니다. 

 

(왼) Inference, (우) test.py

 

유니티 Barraucda 셋팅

 

이 부분에서 RenderTexture를 알고 있으면 궁극적으로 Barracuda 사용에 좋습니다. Rendertexture란 실시간으로 rendering 하는 텍스쳐로 유니티 (RT3D) 및 Barracuda로 실시간 렌더링을 할 경우 보다 시너지를 낼 수 있습니다.

 

본 포스팅에서는 별도로 이 부분을 짚고 넘어가지 않겠습니다.

추가적으로 설치의 경우 이전 포스트에서 다루었기때문에 별도로 다루지 않았습니다.

 

다음과 같이 빈 화면에 Onnx 파일을 넣어주었습니다.

 

빈 오브젝트를 만들고 다음과 같이 코드를 작성하였습니다.

 

다음과 같이 전체 코드를 작성하였습니다.

 

using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using Unity.Barracuda;
using UnityEngine.UI;

public class Inference : MonoBehaviour
{
    public NNModel AdainModel;
    private Model m_RuntimeModel;
    public Texture2D content;
    public RenderTexture result;
    IWorker worker;
    public GameObject _rawstyle;
    RawImage rawImageTexture;
    
        public enum style_list
    {
        antimonocromatismo,
        asheville,
        brushstrokes,
        contrast_of_forms
    }
    public style_list set_style = style_list.antimonocromatismo;
    style_list prev_condition;


    Dictionary<string, Tensor> Inputs = new Dictionary<string, Tensor>();
    // Start is called before the first frame update

    void Set()
    {
        m_RuntimeModel = ModelLoader.Load(AdainModel);
        rawImageTexture = _rawstyle.GetComponent<RawImage>();
        set_style = style_list.antimonocromatismo;
        excute();
    }

    void excute()
    {
        var worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, m_RuntimeModel);
        Inputs = new Dictionary<string, Tensor>();
        Tensor t_content = new Tensor(content, 3);
        Tensor t_style = new Tensor(rawImageTexture.texture, 3);
        Inputs.Add("content", t_content);
        Inputs.Add("style", t_style);

        worker.Execute(Inputs);

        Tensor output = worker.PeekOutput("output");
        output.ToRenderTexture(result);

        t_content.Dispose();
        t_style.Dispose();
        output.Dispose();
    }
    void StyleChanged()
    {    
        Texture2D antimonocromatismo = Resources.Load<Texture2D>("antimonocromatismo");
        Texture2D asheville = Resources.Load<Texture2D>("asheville");
        Texture2D brushstrokes = Resources.Load<Texture2D>("brushstrokes");
        Texture2D contrast_of_forms = Resources.Load<Texture2D>("contrast_of_forms");

        if (set_style == prev_condition)
        {
            Debug.Log("Not updated!");
            return;
        }
        else if (set_style == style_list.antimonocromatismo)
        {
            rawImageTexture.texture = antimonocromatismo;
            prev_condition = style_list.antimonocromatismo;
            excute();
            Debug.Log("Updated!");
        }
        else if (set_style == style_list.asheville)
        {
            rawImageTexture.texture = asheville;
            prev_condition = style_list.asheville;
            excute();
            Debug.Log("Updated!");
        }
        else if (set_style == style_list.brushstrokes)
        {
            rawImageTexture.texture = brushstrokes;
            prev_condition = style_list.brushstrokes;
            excute();
            Debug.Log("Updated!");
        }
        else if (set_style == style_list.contrast_of_forms)
        {
            rawImageTexture.texture = contrast_of_forms;
            prev_condition = style_list.contrast_of_forms;
            excute();
            Debug.Log("Updated!");
        }
    }

    void Awake()
    {
        Set();
    }

    void Update()
    {
        StyleChanged();
    }

}

 

코드 리뷰

 

코드와 관련된 부분은 원문의 문서를 읽는 것을 강력하게 추천드립니다.Barracuda의 getting started 문서가 잘 정리되어있습니다.

GitHub - Unity-Technologies/barracuda-release

 

 

public NNModel AdainModel;
private Model m_RuntimeModel;

다음과 같이 NNModel를 선언하고 해당 값에 onnx 파일인 adain을 드래그앤 드롭합니다.

이 후 m_RuntimeModel은 NNmodel을 불러오기 위해 사용합니다. 

 

AdaIN 모델을 불러오기 위해서는 2개의 Input Texture가 필요합니다. 각각의 Texture를 선언해줍니다. 이 후 Barracuda는 유니티의 실시간 3D 특성상 output을 RenderTexture로 제한합니다. 따라서 output을 담아 줄 RenderTexture도 같이 만들어 주겠습니다.

 

public Texture2D content;
public Texture2D style;
public RenderTexture resuilt;

 

이 후 Worker를 선택할 수 있는데 이는 inference를 할 때 몇가지 옵션을 설정할 수 있습니다. 크게 CPU와 GPU로 나뉘고 각각의 설정에서 3개의 세부 항목이 있습니다. 관련 내용은 Iworker Interface에서 확인하실 수 있습니다.

GitHub - Unity-Technologies/barracuda-release

 

저는 ComputePrecompiled를 사용해서 GPU를 사용할 수 있도록 설정하였습니다. 이 셋팅이 가장 효율적인 방법으로 문서에 나와 있습니다.

 

var worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, m_RuntimeModel);

 

이 후 dictionary로 input 값을 담을 수 있도록 설정하고 각각의 Input을 담아 줍니다.

new Tensor 부분의 (content, 3) 부분은 content 이미지를 3채널로 가져온다는 의미입니다.

 

Inputs = new Dictionary<string, Tensor>();
Tensor t_content = new Tensor(content, 3);
Tensor t_style = new Tensor(style, 3);
Inputs.Add("content", t_content);
Inputs.Add("style", t_style);

 

worker를 run해서 inference를 실행시켜주고 그 값을 output tensor에 담을 수 있도록 설정하였습니다. PeekOutput은 한번 실행할 때만 그 값을 가지고 있습니다. tensor의 이미지 변환은 ToRenderTexture로만 작동합니다. 따라서 다음과 같이 result RenderTexture에 그 값을 담을 수 있도록 하였습니다.

 

worker.Execute(Inputs);
Tensor output = worker.PeekOutput("output");
output.ToRenderTexture(result);

 

여기서 유의미한 부분이 PeekOutput을 사용하면 더 이상의 메모리 할당을 진행하지 않습니다. 따라서 다음 worker.Excute()를 실행할 때 tensor는 그 값을 소실하게 됩니다. 이미 ToRenderTexture로 저장을 하였으니 Dispose()를 합니다.

 

관련해서 worker.CopyOutPut으로 excute를 실행해도 그 값을 가지고 있을 수 있는 방법이 있었는데 이 경우도 Dispose()를 사용하는 것을 권장하였습니다

 

Gitsync

 

무슨 문제인지 모르겠는데 생각보다 메모리를 엄청 잡아먹네요. 재부팅을 하니 거의 사라졌지만 필요시에만 (style 이미지가 수정될 때만 excute()를 하도록 전체 코드를 수정하였으니 참고하시길 바랍니다.

 

전체 내용은 다음 레포에서 확인하실 수 있습니다. (pnltoen/Barracuda_AdaIN)

GitHub - pnltoen/Barracuda_AdaIN