Docker기반의 RESTful api로 pytorch model 배포하기[3-5]

1. Flask RESTful을 이용하여, HTTP로 호출할 수 있는 RESTful api 만들기
2. 학습한 모델의 weight와 graph를 하나의 파일로 만들기(torchscript)
3. 보안을 위해, 2번에서 만든 torchscript파일을 encryption하기
4. 보안을 위해, cython을 이용하여 python script를 library형태로 만들기
5. Docker image 만들고 배포하기

 

지난 글까지는 RESTful api를 만들고 pytorch를 이용하여 학습한 모델을 이용하기 위해서, 이를 torchscript 형태로 변환하였다.

아래와 같이 torchscript로 변환된 모델을 save&load 할 수 있는데

''' save '''
traced_net.save(traced_path)
''' load '''
net = torch.jit.load(traced_path)

만약 배포 환경이 사내에 설치된 server와 같이 보안이 어느정도 보장된 환경이라면 무관할수도 있지만, 위와 같은 방법으로 load된 model은 누구나 그 구조를 볼 수 있다.

즉, 개인의 특별한 아이디어가 들어가 있다면 그 구조를 숨기고 싶을 것이다.

 

그래서 이번 글에서는 python환경에서 손쉽게 사용할 수 있는 encryption library를 이용하여, 이전 글에서 만든 torchscript 파일을 암호화하여 model의 구조를 아무나 볼 수 없도록 만들 것이다.

 

필자는 암호화와 관련된 지식은 가지고 있지 않으므로, 이번 글에서는 주로 source code의 copy&paste가 주된 내용이 될 것이다.

 

암호화를 위한 class의 source code는 아래와 같다.

 

import struct, hashlib
import os
from Crypto.Cipher import AES

class aes_encryptor(object):
    def __init__(self, password):
        self.chunksize = 65536
        self.password = password
        ''' generate key '''
        self.key = hashlib.sha256(self.password.encode('utf-8')).digest()
        ''' initial vector string '''
        self.iv = 'initialvector123'.encode('utf-8')

    def encryption(self, out_dir, file_path, remove_file=True):
        encrypt_out = os.path.join(out_dir, os.path.basename(file_path).replace('.pt', '.enc'))
        encryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
        filesize = os.path.getsize(file_path)
        with open(file_path, 'rb') as infile:
            with open(encrypt_out, 'wb') as outfile:
                outfile.write(struct.pack('<Q', filesize))
                outfile.write(self.iv)
                while True:
                    chunk = infile.read(self.chunksize)
                    if len(chunk) == 0:
                        break
                    elif len(chunk) % 16 != 0:
                        chunk += ' '.encode('utf-8') * (16 - len(chunk) % 16)
                    outfile.write(encryptor.encrypt(chunk))

        if remove_file is True:
            os.remove(file_path)
        print('Encryption is done.')
        print('Password is ')
        print(self.password)

    def decryption(self, in_filename, out_filename):
        chunksize = 24 * 1024
        ''' generate key '''
        key = hashlib.sha256(self.password.encode('utf-8')).digest()
        with open(in_filename, 'rb') as infile:
            origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
            ''' read initial vector '''
            iv = infile.read(16)
            decryptor = AES.new(key, AES.MODE_CBC, iv)
            with open(out_filename, 'wb') as outfile:
                while True:
                    chunk = infile.read(chunksize)
                    if len(chunk) == 0:
                        break
                    outfile.write(decryptor.decrypt(chunk))
                outfile.truncate(origsize)

 

암호화 알고리즘은 AES라는 알고리즘을 사용하였으며, AES 알고리즘 중에서도 CBC라는 방법을 이용하였다.

자세한 알고리즘 내용은 따로 검색해보시길 바랍니다.

 

간단하게 설명을 하면, encryption function을 보면 AES.new function에서 key와 mode 그리고 initial vector를 매개변수로 받아서 encryptor를 만들고, 해당 encryptor를 이용하여 암호화하고자 하는 파일을 binary chunk로 나누어서 해당 chunk를 암호화하고 write한다.

여기서 key는 주어진 password를 hashlib를 이용하여 key를 생성하였다.

descryption도 encryption과 크게 다르지 않다.

 

해당 source code를 이용하여, torchscript로 만든 파일을 save&load하는 흐름은 아래와 같다.

 

''' make a torchscript '''

''' crate instance '''
aes = aes_encryptor(password='mypassword')
''' encryption '''
aes.encryption(out_dir='tmp', file_path='torchscript.pt', remove_file=True)
''' decryption '''
aes.decryption(in_filename=os.path.join('tmp', 'net.enc'), out_filename='tmp.pt')

 

위의 code를 이용하여 torchscript 파일을 암호화하여 model의 구조를 보호할 수 있다.