SinLU를 LibTorch로 구현

원본 논문#

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

SinLU 구조#

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

1
import torch
2
import torch.nn as nn
3

4
class SinLU(nn.Module):
5
    def __init__(self):
6
        super(SinLU,self).__init__()
7
        self.a = nn.Parameter(torch.ones(1))
8
        self.b = nn.Parameter(torch.ones(1))
9
    def forward(self,x):
10
        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

1
.
2
├── build
3
│   └── filelocation
4
├── build.py
5
├── CMakeLists.txt
6
├── main.cpp
7
├── model.h
8
├── README.md
9
└── sinlu.h

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

1. sinlu.h#

1
#pragma once
2
#include <torch/torch.h>
3
#include <torch/nn.h>
4

5
namespace sinlu{
6
typedef struct SinLUImpl : torch::nn::Module {
7
    SinLUImpl(int64_t N=1)
8
    {
9
        a = register_parameter("a", torch::ones(N), true);
10
        b = register_parameter("b", torch::ones(N), true);
11
    }
12
    torch::Tensor forward(torch::Tensor x) {
13
        return torch::sigmoid(x)*(x+a*torch::sin(b*x));
14
        }
15
    torch::Tensor a, b;
16
    }SinLUImpl;
17
    TORCH_MODULE(SinLU);
18
}

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

2. model.h#

1
#pragma once
2
#include "sinlu.h"
3

4
namespace model{
5

6
typedef struct ModelImpl : torch::nn::Module {
7
    ModelImpl(int kSize):
8
        linear1(torch::nn::LinearOptions(kSize, 5)),
9
        sinlu1(1),
10
        linear2(torch::nn::LinearOptions(5, 2))
11
    {
12
    register_module("linear1", linear1);
13
    register_module("sinlu1", sinlu1);
14
    register_module("linear2", linear2);
15
    }
16

17
    torch::Tensor forward(torch::Tensor x) {
18
        x = linear1(x);
19
        x = sinlu1(x);
20
        x = linear2(x);
21
        return torch::sigmoid(x);
22
    }
23

24
    torch::nn::Linear linear1, linear2;
25
    sinlu::SinLU sinlu1;
26
    }ModelImpl;
27
    TORCH_MODULE(Model);
28
}

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

3. main.cpp#

1
#include "sinlu.h"
2
#include "model.h"
3

4
int main() {
5

6
    int i;
7
    torch::manual_seed(1234);
8
    model::Model model(10);
9

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

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

24
    return 0;
25
}

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