mobile_net.py 4.45 KB
"""
    ATTENTION을 적용한 MOBILE NET
    @FUNCTION load_data : pickle 데이터를 로딩하는 함수
    @FUNCTION Mobile_net : 모바일 넷 모델 함수
    @FUNCITON predict : 모델 예측하는 함수
"""
import pickle
import numpy as np
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import keras
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense, Conv2D, Conv3D, DepthwiseConv2D, SeparableConv2D, Conv3DTranspose
from tensorflow.keras.layers import Flatten, MaxPool2D, AvgPool2D, GlobalAvgPool2D, UpSampling2D, BatchNormalization
from tensorflow.keras.layers import Concatenate, Add, Dropout, ReLU, Lambda, Activation, LeakyReLU, PReLU

from attention_module import cbam_block

"""
    데이터 로드
    @brief : load Inputs and Targets from pickle data 
    @param data_path(str) : path to pickle file containing data
    @return X(ndarray) : Inputs
    @return y(ndarray) : Targets
"""
def load_data():
    X = pickle.load(open("X_final2.pickle", "rb"))
    y = pickle.load(open("y_final2.pickle", "rb"))
    
    X = X/225.0
    
    return X, y
    
"""
    mobile net 구현
    @brief : Mobile net with Convolution Block Attention Module(CBAM)
            I used cbam at last of convolution when I used it at every conv block,
            the result was worse.
    @return model : Mobile Net Model
"""
def mobile_net(input_shape):
    
    def mobile_net_block(x, f, s=1):
        x = DepthwiseConv2D(3, strides=s, padding='same')(x)
        x = BatchNormalization()(x)
        x = ReLU()(x)
        
        
        x = Conv2D(f, 1, strides=1, padding='same')(x)  
        x = BatchNormalization()(x)
        x = ReLU()(x)
        
        
        return x
    
    input = Input(input_shape)
    
    x = Conv2D(32, 3, strides=2, padding='same')(input)
    x = BatchNormalization()(x)
    x = ReLU()(x)
    
    x = mobile_net_block(x, 64)
    x = mobile_net_block(x, 128, 2)
    x = mobile_net_block(x, 128)
    
    x = mobile_net_block(x, 256, 2)
    x = mobile_net_block(x, 256)
    
    x = mobile_net_block(x, 512, 2)
    for _ in range(5):
        x = mobile_net_block(x, 512)
        
    x = mobile_net_block(x, 1024, 2)
    x = mobile_net_block(x, 1024)
    x = cbam_block(x)
    
    x = GlobalAvgPool2D()(x)
    
    
    output = Dense(3, activation='softmax')(x)
    
    model = Model(input, output)
    return model

"""
    학습된 모델로 예측
    @brief : predict data from trained mobile net model
    @param model : Trained classifier
    @param X : Input data
    @param y(int): Target
"""
def predict(model, X, y):
    # sample의 입력 데이터에 차원 추가
    X = X[np.newaxis, ...]
    
    prediction = model.predict(X)
    
    # argmax를 사용해서 index의 최대 값을 얻음
    predicted_index = np.argmax(prediction, axis=1)
    
    print("Target: {}, Predicted label: {}".format(y, predicted_index))


# 메인 함수
if __name__ == "__main__":
    
    # load data and split to X_train and y_train
    X_train, y_train = load_data()
    
    # create network
    K.clear_session()
    input_shape = (X_train.shape[1], X_train.shape[2], 3)
    model= mobile_net(input_shape)
    
    # compile model
    optimiser = keras.optimizers.Adam(learning_rate=0.0001)
    model.compile(optimizer=optimiser,
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    
    model.summary()
    
    # train model
    hist = model.fit(X_train, y_train, validation_split=0.1, batch_size=32, epochs=30)
    
    fig, loss_ax = plt.subplots()
    
    acc_ax = loss_ax.twinx()
    
    loss_ax.plot(hist.history['loss'], 'y', label='train loss')
    loss_ax.plot(hist.history['val_loss'], label='validation loss')
    
    acc_ax.plot(hist.history['acc'], 'b', label='train_acc')
    acc_ax.plot(hist.history['val_acc'], 'g', label='validation_acc')
    
    loss_ax.set_xlabel('epoch')
    loss_ax.set_ylabel('loss')
    acc_ax.set_ylabel('accuracy')
    
    loss_ax.legend(loc='upper left')
    loss_ax.legend(loc='lower left')
    
    # evalute model
    test_loss, test_acc = model.evaluate(X_train, y_train, verbose=2)
    print('\nTest accuracy:', test_acc)
    
    X_to_predict = X_train[100]
    y_to_predict = y_train[100]
    
    # predict sample
    predict(model, X_to_predict, y_to_predict)
    
    model.save('C:/Users/nokh9/Desktop/mobile_net_final2.h5')