SinLU를 LibTorch로 구현

원본 논문

SinLU: Sinu-sigmoidal linear unit. Mathematics, 2022.

SinLU 구조

Github를 참고하면 SinLU의 코드는 다음과 같이 정의되어 있다.

import torch
import torch.nn as nn

class SinLU(nn.Module):
    def __init__(self):
        super(SinLU,self).__init__()
        self.a = nn.Parameter(torch.ones(1))
        self.b = nn.Parameter(torch.ones(1))
    def forward(self,x):
        return torch.sigmoid(x)*(x+self.a*torch.sin(self.b*x))

간단한 코드이며, nn.Parameter를 통해 해당 tensor의 값을 파라미터로 등록하여 모델이 업데이트 (gradient) 될때 해당 값도 업데이트된다.

프로젝트 환경 및 구조

LibTorch Release CPU version (1.12.0)
Visual Studio build tools 2019
Cmake

.
├── build
│   └── filelocation
├── build.py
├── CMakeLists.txt
├── main.cpp
├── model.h
├── README.md
└── sinlu.h

먼저 sinlu.h부터 살펴보자.

1. sinlu.h

#pragma once
#include <torch/torch.h>
#include <torch/nn.h>

namespace sinlu{
typedef struct SinLUImpl : torch::nn::Module {
    SinLUImpl(int64_t N=1)
    {
        a = register_parameter("a", torch::ones(N), true);
        b = register_parameter("b", torch::ones(N), true);
    }
    torch::Tensor forward(torch::Tensor x) {
        return torch::sigmoid(x)*(x+a*torch::sin(b*x));
        }
    torch::Tensor a, b;
    }SinLUImpl;
    TORCH_MODULE(SinLU);
}

메소드의 구현체인 Impl를 붙여SinLUImpl를 선언한다. nn.Parameter와 같은 역할을 하는게 register_parameter이다. TORCH_MODULE를 이용하여 SinLU을 사용할 수 있도록 한다.

2. model.h

#pragma once
#include "sinlu.h"

namespace model{

typedef struct ModelImpl : torch::nn::Module {
    ModelImpl(int kSize):
        linear1(torch::nn::LinearOptions(kSize, 5)),
        sinlu1(1),
        linear2(torch::nn::LinearOptions(5, 2))
    {
    register_module("linear1", linear1);
    register_module("sinlu1", sinlu1);
    register_module("linear2", linear2);
    }

    torch::Tensor forward(torch::Tensor x) {
        x = linear1(x);
        x = sinlu1(x);
        x = linear2(x);
        return torch::sigmoid(x);
    }

    torch::nn::Linear linear1, linear2;
    sinlu::SinLU sinlu1;
    }ModelImpl;
    TORCH_MODULE(Model);
}

ModelImpl를 선언한다. sinlu.h에서 만든 SinLU를 이용하여 모델을 만든다. 각 레이어의 출력은 여기서는 하드코딩으로 처리했지만 #define으로 설정하는 편이 깔끔할 것이다.

3. main.cpp

#include "sinlu.h"
#include "model.h"

int main() {

    int i;
    torch::manual_seed(1234);
    model::Model model(10);

    torch::optim::Adam optimizer(
        model->parameters());

    torch::Tensor feature = torch::rand({1, 10});
    torch::Tensor target = torch::rand({1, 2});
    for (i = 0; i < 100; i++){
        model->zero_grad();
        torch::Tensor out = model->forward(feature);
        torch::Tensor loss = torch::binary_cross_entropy(out, target);
        loss.backward();
        optimizer.step();
        std::printf("\r Epoch : %d, Loss : %f\n", i, loss.item<float>());
    }

    return 0;
}

입력되는 값들은 랜덤으로 처리하였고, epoch 횟수를 i로 처리하였다. 메소드는 PyTorch와 같은 이름을 가진 메소드가 있어 그대로 사용하면 된다.