[tensorflow] 모든 GPU 사용하여 모델 predict

---

Tensorflow 모델 트레이닝 방법 중에 모든 gpu를 태워서 '학습(training)' 하는 방법에 대한 내용은 정리가 많이 되어있는 편입니다. 

하지만 기존의 만들어진 모델을 각 GPU에 모두 태워서 '예측(predict)' 하는 방법은 거의 정리된 내용이 없어서 제가 겪은 뻘짓과 성공 경험을 정리해두려고 합니다. 

---

분석 환경

AWS GPU instance 
Tesla M60 x 4 
Memory : 480 G
CPUs : 4 CPU
N of Cores per CPU : 16 cores 

목적 

• 매일 쌓이는 상품코드에 대한 상품 이미지(url)를 2048 길이로 임베딩하여 저장. 
• 많이 등록될 때에는 약 5000개의 상품이 등록되는 경우가 있어, 이를 자동화하여 배치 작업으로 수행할 수 있도록 함. 
• 기존의 작업은 Jupyter lab에서 한 커널 당 하나의 GPU를 담당하여 수동으로 임베딩 작업을 거쳤기 때문에 아주 번거롭고 비효율적인 작업.

GPU 할당 방법 예시.

os.environ["CUDA_VISIBLE_DEVICES"] = '3' # 4번째 GPU를 담당하는 Jupyter 커널. 

 

작업을 수행하던 중, 각각 임베딩을 수행하던 커널을 인스턴스화하여 실행하면 어떨까? 생각하였습니다. 

Python - 클래스와 인스턴스를 대충 한번 훑어봤던 기억을 더듬으며 아래와 같이 tensorflow, keras 의 xception 임베딩 모델을 GPU를 지정하여 구동할 수 있는 클래스를 만들어봤습니다. 

import os
import gc
import pickle
import gzip
import math

from PIL import Image # Pillow, OpenCV, PIL 
from io import BytesIO
from socket import timeout
from urllib.request import urlopen
import urllib.parse
import urllib.request

import cx_Oracle
import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.applications.xception import Xception
from tensorflow.python.keras.layers import GlobalMaxPooling2D
from tensorflow.python.keras.applications.xception import preprocess_input, decode_predictions

class Embed_predictor():
    def __init__(self, cuda_divice_num = 1):
        
        os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_divice_num)
        self.vector_shape = (813, 640, 3)
        self.model = self.make_embedding_model(self.vector_shape)
        
    def make_embedding_model(self, vector_shape):
        base_model = Xception(weights = 'imagenet',    # 이미지넷으로 트레이닝 된 모델(Xception)을 불러옴
                              include_top = False,     # 상단의 Fully-connected layer를 포함할 것인지 아닌지 결정 
                              input_shape = vector_shape)

        base_model.trainable = False                   # Feature extraction 파라미터들은 이미지넷으로 학습된 값들을 
                                                       # 그대로 사용할 것이기 때문에 trainable의 속성을 False로 함

        # Add Layer Embedding
        xception_model = tf.keras.Sequential([
                                              base_model,
                                              GlobalMaxPooling2D()
                                             ])
        
        return xception_model
    
    
    def load_img_from_url(self, url):
        resp = urlopen(url)
        im_src = Image.open(BytesIO(resp.read()))
        mode = im_src.mode

        if mode != 'RGB':
            im = im_src.convert('RGB')
        else:
            im = im_src

        img_vector = np.array(im, dtype=np.uint8)
        return img_vector, mode
        
    def get_embedding(self, img_name):
        img_vector, mode = self.load_img_from_url(img_name)

        try:
            img_vector = np.expand_dims(img_vector, axis=0)
            preprocess_image = preprocess_input(img_vector) 
            return self.model.predict(preprocess_image).reshape(-1)

        except Exception as e:
            return np.zeros(2048)
            print("e: ", e)        

    def return_df(self, img_shape_df):
        emb_img_list = []
        for img_url in img_shape_df['IMAGE']:
            img_emb = self.get_embedding(img_url)
            emb_img_list.append(img_emb)
            
        lf_img_list = pd.DataFrame(emb_img_list)
        img_emb_df = pd.concat([img_shape_df[['PROD_CD', 'IMAGE']], lf_img_list], axis=1)
            
        return img_emb_df

 

해당 클래스는 인스턴스 생성 시, GPU 번호와 xception 모델을 생성할 수 있도록 하였습니다. 각 인스턴스당 GPU를 파이썬 분산처리 라이브러리 multiprocessing를 활용하여 각각 프로세스 당 하나씩 할당하도록 하였습니다.  

Why? - Tensorflow 임베딩 모델 Xception 모델을 만드는 동시에 GPU 메모리를 할당하기 때문에, 각 인스턴스 당 GPU를 지정한 이후에 모델 생성을 수행해야 하였습니다. 

 

[인스턴스 선언(초기화) 시 GPU 할당 및 모델 생성]

def __init__(self, cuda_divice_num = 1):

    os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_divice_num) # GPU 할당
    self.vector_shape = (813, 640, 3) 
    self.model = self.make_embedding_model(self.vector_shape) # xception 모델 생성

 

해당 클래스를 사용하는 방법은 다음과 같습니다. 

테이블 예시.

 

상품코드 - 이미지 URL

A9030 - https://imageurl.com?A9030

B1203  - https://imageurl.com?B1203 

# use multi gpu
from predictor import Embed_predictor
import pandas as pd
import numpy as np 

def predict_GPUs(num, df):
    predictor = Embed_predictor(num)
    return predictor.return_df(df.reset_index())

def paralle_GPUs(df, func, n_cores = 4):
    df_split = np.array_split(df, n_cores)    
    a_args = [0,1,2,3] # GPU list 
    
    with Pool(4) as pool:
        lst = pool.starmap(func, zip(a_args, df_split))
    
    return lst   
    
# 이미지 url이 포함된 상품테이블 입력
# paralle_GPUs(테이블, 함수, 프로세서 수) 
result = paralle_GPUs(img_shape_df, predict_GPUs, 4)

 

해당 내용은 아래 깃허브에 정리해두었습니다. 

https://github.com/Yoon-Alex/multi_GPUs

GitHub - Yoon-Alex/multi_GPUs: 각 인스턴스 당 GPU를 할당, 분산 처리(predict)을 할 수 있도록 구성한 작업.