Skip to content
Snippets Groups Projects
vgg16_imagenet.py 6.50 KiB
import os
import glob
import random

import scipy
import scipy.io
import cv2
import numpy as np

import tensorflow as tf
import keras
from keras.models import Sequential, Model
from keras.layers import *
from keras.utils import to_categorical
from keras.applications.vgg16 import VGG16, preprocess_input
from keras import backend as K

from frontend.approxhpvm_translator import translate_to_approxhpvm
from frontend.weight_utils import dumpCalibrationData2


np.random.seed(2020)

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

K.set_image_data_format('channels_first')

data_format = 'channels_first'


IMAGENET_DIR = '/home/nz11/ILSVRC2012/'
OUTPUT_DIR = 'data/vgg16_imagenet_tune/'

NUM_CLASSES = 200
IMAGES_PER_CLASS = 50
# VAL_SIZE = 100



def get_vgg16_nchw_keras():
    img_input = Input(shape=(3, 224, 224))
    
    # Block 1
    x = Conv2D(64, (3, 3),
                      padding='same',
                      data_format=data_format)(img_input)
    x = Activation('relu')(x)
    x = Conv2D(64, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), data_format=data_format)(x)

    # Block 2
    x = Conv2D(128, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(128, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), data_format=data_format)(x)

    # Block 3
    x = Conv2D(256, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(256, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(256, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), data_format=data_format)(x)

    # Block 4
    x = Conv2D(512, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(512, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(512, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), data_format=data_format)(x)

    # Block 5
    x = Conv2D(512, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(512, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = Conv2D(512, (3, 3),
                      padding='same',
                      data_format=data_format)(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), data_format=data_format)(x)

    x = Flatten(data_format='channels_last')(x)
    
    x = Dense(4096)(x)
    x = Activation('relu')(x)
    x = Dense(4096)(x)   
    x = Activation('relu')(x)
    x = Dense(1000)(x)
    x = Activation('softmax')(x)

    model_nchw = Model(img_input, x)
    
    
    model = VGG16()

    j = 0
    for i in range(len(model_nchw.layers)):
        if 'padding' in model_nchw.layers[i].name or 'activation' in model_nchw.layers[i].name:
            continue
        try:
            model_nchw.layers[i].set_weights(model.layers[j].get_weights())
        except:
            print (i, model_nchw.layers[i], 'skipped')
        j += 1

    
    return model_nchw



def load_image(x):
    
    image = cv2.imread(x)
        
    height, width, _ = image.shape
    new_height = height * 256 // min(image.shape[:2])
    new_width = width * 256 // min(image.shape[:2])
    image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
    
    height, width, _ = image.shape
    startx = width // 2 - (224 // 2)
    starty = height // 2 - (224 // 2)
    image = image[starty:starty + 224, startx:startx + 224]
    
    image = image[:, :, ::-1]
    image = np.transpose(image, (2, 0, 1))
    image = preprocess_input(image.astype(np.float32), data_format=data_format)
        
    return image.astype(np.float32)


meta = scipy.io.loadmat(IMAGENET_DIR + 'ILSVRC2012_devkit_t12/data/meta.mat')
original_idx_to_synset = {}
synset_to_name = {}

for i in range(1000):
    ilsvrc2012_id = int(meta['synsets'][i,0][0][0][0])
    synset = meta['synsets'][i,0][1][0]
    name = meta['synsets'][i,0][2][0]
    original_idx_to_synset[ilsvrc2012_id] = synset
    synset_to_name[synset] = name

synset_to_keras_idx = {}
keras_idx_to_name = {}
f = open(IMAGENET_DIR + 'ILSVRC2012_devkit_t12/data/synset_words.txt', 'r')
c = 0
for line in f:
    parts = line.split(' ')
    synset_to_keras_idx[parts[0]] = c
    keras_idx_to_name[c] = ' '.join(parts[1:])
    c += 1
f.close()




model = get_vgg16_nchw_keras()
    
X_tune, X_test = [], []
y_tune, y_true = [], []
    
classes = glob.glob(IMAGENET_DIR + 'val/*')

for c in np.random.permutation(len(classes))[:NUM_CLASSES]:
    x = glob.glob(classes[c] + '/*')
    x = np.array(x)
            
    idx = np.random.permutation(len(x))
    idx = idx[:max(len(idx), IMAGES_PER_CLASS)]
    
    synset = classes[c].split('/')[-1]
    images = list(map(lambda x : load_image(x), x[idx]))
    labels = [synset_to_keras_idx[synset]] * len(x[idx])
    
    X_test += images[:IMAGES_PER_CLASS // 2]
    y_true += labels[:IMAGES_PER_CLASS // 2]
    
    X_tune += images[IMAGES_PER_CLASS // 2:]
    y_tune += labels[IMAGES_PER_CLASS // 2:]
    
    
X_test = np.array(X_test)
y_true = np.array(y_true)
X_tune = np.array(X_tune)
y_tune = np.array(y_tune)

print ('tune size', len(X_tune))
print ('test size', len(X_test))



translate_to_approxhpvm(model, OUTPUT_DIR, X_tune, y_tune, 1000)

# # dumpCalibrationData2(OUTPUT_DIR + 'test_input_10K.bin', X_test, OUTPUT_DIR + 'test_labels_10K.bin', y_true)
# dumpCalibrationData2(OUTPUT_DIR + 'tune_input.bin', X_tune, OUTPUT_DIR + 'tune_labels.bin', y_tune)
# dumpCalibrationData2(OUTPUT_DIR + 'test_input.bin', X_test, OUTPUT_DIR + 'test_labels.bin', y_true)



pred = np.argmax(model.predict(X_test), axis=1)
print ('val accuracy', np.sum(pred == y_true.ravel()) / len(X_test))
    
# pred = np.argmax(model.predict(X_tune), axis=1)
# print ('val accuracy', np.sum(pred == y_tune.ravel()) / len(X_tune))