train_task_classifyapp.py

# NCC: Neural Code Comprehension
# https://github.com/spcl/ncc
# Copyright 2018 ETH Zurich
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
# following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
# disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
# disclaimer in the documentation and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
# products derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# ==============================================================================
"""Training workflow for app classification"""
from labm8 import fs
import task_utils
import rgx_utils as rgx
import pickle
from sklearn.metrics import confusion_matrix
from sklearn.utils import resample
import os
import sys
import json
import numpy as np
import math
import struct
import tensorflow as tf
from tensorflow.keras.callbacks import Callback
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.layers import Input, LSTM, Dense, Embedding, BatchNormalization
from tensorflow.keras.models import Model
from absl import flags
import matplotlib.pyplot as plt
from scipy.special import softmax

# Parameters of classifyapp
flags.DEFINE_string('input_data', 'task/classifyapp_lifted', 'path to input data')
flags.DEFINE_string('folder_results', 'task/classifyapp_lifted', 'path to write model weights and predictions')
flags.DEFINE_integer('num_epochs', 50, 'number of training epochs')
flags.DEFINE_integer('batch_size', 32, 'training batch size')
flags.DEFINE_integer('dense_layer_size', 32, 'dense layer size')
flags.DEFINE_integer('trsamples', 398, 'number of training samples per class, if 0 use all')
flags.DEFINE_integer('vasamples', 0, 'number of validation samples per class, if 0 use all')
flags.DEFINE_integer('tesamples', 0, 'number of test samples per class, if 0 use all')
flags.DEFINE_integer('maxlen', 1100, 'max length of sequences, all sequences padded or cuted to this number, '
                                     'if 0 specified, then compute it dynamically')
flags.DEFINE_string('model_name', "NCC_classifyapp_lifted", 'name of model to train or use for predictions')
flags.DEFINE_integer('save_every', 10, 'save checkpoint every N batches')
flags.DEFINE_integer('ring_size', 5, 'checkpoint ring buffer length')
flags.DEFINE_bool('print_summary', False, 'print summary of Keras model')
flags.DEFINE_string('mode', 'predict', 'train/test/predict')
flags.DEFINE_string('json_file', 'NCC_classifyapp_lifted.json', 'filename to store predictions')
flags.DEFINE_integer('topk', 3, 'how match labels to see per sample for computing accuracy')
flags.DEFINE_integer('seed', 204, 'initializer for pseudorandom generators')
flags.DEFINE_integer('num_classes', 104, 'number of classes')
FLAGS = flags.FLAGS
FLAGS(sys.argv)


def get_onehot(y, num_classes):
    """
    y is a vector of numbers in range (1, number of classes)
    """
    hot = np.zeros((len(y), num_classes), dtype=np.int32)
    for i, c in enumerate(y):
        hot[i][c - 1] = 1
    return hot


def encode_srcs(input_files, dataset_name, unk_index):
    """
    encode source code for learning
    data_folder: folder from which to read input files
    input_files: list of strings of file names
    """

    num_files = len(input_files)
    num_unks = 0
    seq_lengths = list()
    unk_index = str(unk_index)

    print('\n--- Preparing to read', num_files, 'input files for', dataset_name, 'data set')
    seqs = list()
    for i, file in enumerate(input_files):
        if i % 10000 == 0:
            print('\tRead', i, 'files')
        with open(file, 'rb') as f:
            full_seq = f.read()
        seq = list()
        for j in range(0, len(full_seq), 4):  # read 4 bytes at a time
            seq.append(struct.unpack('I', full_seq[j: j + 4])[0])
        assert len(seq) > 0, 'Found empty file: ' + file
        num_unks += seq.count(unk_index)
        seq_lengths.append(len(seq))
        seqs.append([int(s) for s in seq])

    maxlen = max(seq_lengths)
    print('\tShortest sequence    : {:>5}'.format(min(seq_lengths)))
    print('\tLongest sequence     : {:>5}'.format(maxlen))
    print('\tMean sequence length : {:>5} (rounded down)'.format(math.floor(np.mean(seq_lengths))))
    print('\t0.75 quantile sequence length : {:>5} (rounded down)'.format(math.floor(np.quantile(seq_lengths, 0.75))))
    print('\tNumber of \'UNK\'      : {:>5}'.format(num_unks))
    print('\tPercentage of \'UNK\'  : {:>5.2} (% among all stmts)'.format((num_unks * 100) / sum(seq_lengths)))
    print('\t\'UNK\' index          : {:>5}'.format(unk_index))
    return seqs, maxlen


def pad_src(seqs, maxlen, unk_index):
    encoded = pad_sequences(seqs, maxlen=maxlen, value=unk_index)
    return encoded


def return_classes_and_probs(preds, k):
    classes = np.transpose(np.transpose(preds.argsort())[-k:][::-1]) + 1
    probabilities = np.transpose(np.transpose(np.sort(softmax(preds, axis=-1)))[-k:][::-1])
    return classes, probabilities


def print_stat(y_test, classes, model_name, folder_results):
    topk = classes.shape[1]
    accuracy = np.zeros_like(y_test)
    classes = np.transpose(np.array(classes))
    for i in range(topk):
        accuracy += np.array(classes[i]) == y_test
    print('\nTest top{} accuracy: {:.2f}%'.format(topk, sum(accuracy) * 100.0 / len(accuracy)))

    conf_matr = confusion_matrix(y_test, classes[0])

    fig, ax = plt.subplots()
    values = plt.imshow(conf_matr)
    ax.xaxis.tick_top()
    ax.xaxis.set_label_position('top')

    fig.colorbar(values)
    ax.set_xlabel('Настоящие классы')
    ax.set_ylabel('Предсказанные классы')
    conf_png = os.path.join(folder_results, "models/conf_matr_{}.png".format(model_name))
    plt.savefig(conf_png)


class WeightsSaver(Callback):
    def __init__(self, model, save_every, ring_size, weights_bpath, model_name):
        super().__init__()
        self.model = model
        self.save_every = save_every
        self.ring_size = ring_size
        self.batch = 0
        self.ring = 0
        self.weights_bpath = weights_bpath
        self.model_name = model_name

    def on_batch_end(self, batch, logs={}):
        if self.batch % self.save_every == 0:
            name = os.path.join(self.weights_bpath, '{}_weights{}.h5'.format(self.model_name, self.ring))
            self.model.save_weights(name)
            self.ring = (self.ring + 1) % self.ring_size
        self.batch += 1


class NCC_classifyapp:
    def __init__(self, seed, maxlen, num_classes, dense_layer_size, embedding_matrix,
                 folder_results, model_name, unk_index, topk, save_every, ring_size):
        self.name = model_name
        self.num_classes = num_classes
        self.folder_results = folder_results
        self.unk_index = unk_index
        self.topk = topk
        self.maxlen = maxlen
        self.predictions_path = os.path.join(os.path.join(folder_results, "predictions"),
                                             "{}_top{}.result".format(model_name, topk))
        self.weights_path = os.path.join(os.path.join(self.folder_results, "models"), self.name + '_weights.h5')
        self.seed = seed
        self.save_every = save_every
        self.ring_size = ring_size
        np.random.seed(seed)

        _inp = Input(shape=(maxlen, ), dtype='int32', name='code_in')
        embedding_dim = embedding_matrix.shape[1]
        embedding_layer = Embedding(input_dim=embedding_matrix.shape[0], output_dim=embedding_dim,
                                    weights=[embedding_matrix], input_length=maxlen, trainable=False)
        inp = embedding_layer(_inp)
        x = LSTM(embedding_dim, implementation=1, return_sequences=True, name="lstm_1")(inp)
        x = LSTM(embedding_dim, implementation=1, name="lstm_2")(x)
        x = BatchNormalization()(x)
        x = Dense(dense_layer_size, activation="relu")(x)
        outputs = Dense(num_classes)(x)

        self.model = Model(inputs=_inp, outputs=outputs)
        self.model.compile(optimizer="Adam",
                           loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
        print('\tbuilt Keras model')

    def save(self):
        self.model.save(self.weights_path)

    def restore(self):
        self.model.load_weights(self.weights_path)

    def train(self, sequences: np.array, y_1hot: np.array, sequences_val: np.array, y_1hot_val: np.array,
              verbose: bool, epochs: int, batch_size: int) -> None:
        checkpoint = WeightsSaver(self.model, self.save_every, self.ring_size,
                                  os.path.join(self.folder_results, "models"), self.name)
        try:
            self.model.fit(x=sequences, y=y_1hot, epochs=epochs, batch_size=batch_size, verbose=verbose, shuffle=True,
                           validation_data=(sequences_val, y_1hot_val), callbacks=[checkpoint])

        except KeyboardInterrupt:
            print('training was interrupted, saving weights to file')
            self.model.save_weights(self.weights_path)

    def predict(self, sequences: np.array, batch_size: int) -> np.array:
        p = np.array(self.model.predict(sequences, batch_size=batch_size, verbose=0), dtype=int)
        return return_classes_and_probs(p, self.topk)


def train_model(model, folder_data, folder_results, print_summary, num_epochs, batch_size,
                trsamples, vasamples, tesamples):
    """
    prepare data and train model
    """
    assert os.path.exists(os.path.join(folder_data, 'ir_train')), "Folder not found: " + folder_data + '/ir_train'
    task_utils.llvm_ir_to_trainable(os.path.join(folder_data, 'ir_train'))
    assert os.path.exists(os.path.join(folder_data, 'ir_val')), "Folder not found: " + folder_data + '/ir_val'
    task_utils.llvm_ir_to_trainable(os.path.join(folder_data, 'ir_val'))
    assert os.path.exists(os.path.join(folder_data, 'ir_test')), "Folder not found: " + folder_data + '/ir_test'
    task_utils.llvm_ir_to_trainable(os.path.join(folder_data, 'ir_test'))

    os.makedirs(folder_results, exist_ok=True)
    os.makedirs(os.path.join(folder_results, "models"), exist_ok=True)
    os.makedirs(os.path.join(folder_results, "predictions"), exist_ok=True)

    seed = model.seed
    num_classes = model.num_classes
    X_train, X_val, X_test = list(), list(), list()
    Y = {"train": np.empty(0, dtype=np.int32), "val": np.empty(0, dtype=np.int32), "test": np.empty(0, dtype=np.int32)}
    print('Getting file names for', num_classes, 'classes from folder:', folder_data)

    for i in range(1, num_classes + 1):
        for dataset, nsamples, X in zip(('train', 'val', 'test'), (trsamples, vasamples, tesamples),
                                        (X_train, X_val, X_test)):
            folder = os.path.join(os.path.join(folder_data, 'seq_{}'.format(dataset)), str(i))
            assert os.path.exists(folder), "Folder: " + folder + ' does not exist'
            print('\t{}  : Read file names from folder '.format(dataset), folder)
            seq_files = [os.path.join(folder, f) for f in os.listdir(folder) if f[-4:] == '.rec']
            nsamples = len(seq_files) if nsamples == 0 else nsamples if nsamples <= len(seq_files) else len(seq_files)
            X += resample(seq_files, replace=False, n_samples=nsamples, random_state=seed)
            Y[dataset] = np.concatenate([Y[dataset], np.array([int(i)] * nsamples, dtype=np.int32)])

    unk_index = model.unk_index
    y_train, y_val, y_test = Y["train"], Y["val"], Y["test"]
    del Y
    X_seq_train, maxlen_train = encode_srcs(X_train, 'training', unk_index)
    X_seq_val, maxlen_val = encode_srcs(X_val, 'validation', unk_index)
    X_seq_test, maxlen_test = encode_srcs(X_test, 'testing', unk_index)

    maxlen = max(maxlen_train, maxlen_test, maxlen_val)
    print('Max. sequence length overall:', maxlen)
    maxlen = model.maxlen if model.maxlen else maxlen
    print('Padding sequences to length', maxlen)

    X_seq_train = pad_src(X_seq_train, maxlen, unk_index)
    X_seq_val = pad_src(X_seq_val, maxlen, unk_index)
    X_seq_test = pad_src(X_seq_test, maxlen, unk_index)

    y_1hot_train = get_onehot(y_train, num_classes)
    y_1hot_val = get_onehot(y_val, num_classes)
    del y_train, y_val

    model_name = model.name
    weights_path = model.weights_path
    predictions_path = model.predictions_path

    if fs.exists(predictions_path):
        print("\tFound predictions in", predictions_path, ", skipping...")
        with open(predictions_path, 'rb') as infile:
            classes = np.array(pickle.load(infile))
    else:
        if fs.exists(weights_path):
            print("\n\tFound trained model in", weights_path, ", skipping...")
            model.restore()
        else:
            if print_summary:
                model.model.summary()
            print('\n--- Training model...')
            model.train(X_seq_train, y_1hot_train, X_seq_val, y_1hot_val, True, num_epochs, batch_size)
            model.save()
            print('\tstore model to', weights_path)

        print('\n--- Testing model...')
        classes, probabilities = model.predict(X_seq_test, batch_size)
        del probabilities
        with open(predictions_path, 'wb') as outfile:
            pickle.dump(classes.tolist(), outfile)
        print('\tstore predictions to', predictions_path)

    print_stat(y_test, classes, model_name, folder_results)


def predict_labels(model, json_out):
    """
    predict classes for .ll files in inference/ir_test folder
    """
    unk_index = model.unk_index
    path = 'inference'
    task_utils.llvm_ir_to_trainable(os.path.join(path, 'ir_test'))
    files = [os.path.join(os.path.join(path, 'seq_test'), f) for f in os.listdir(os.path.join(path, 'seq_test')) if
             f[-4:] == '.rec']

    batch_size = len(files)
    X_seq_test, maxlen = encode_srcs(files, 'predict_sample', unk_index)
    print('Max. sequence length overall:', maxlen)
    maxlen = model.maxlen if model.maxlen else maxlen
    print('Padding sequences to length', maxlen)
    X_seq_test = pad_src(X_seq_test, maxlen, unk_index)

    model.model.summary()
    model.restore()
    classes, probabilities = model.predict(X_seq_test, batch_size)

    with open(os.path.join(path, json_out), 'w') as file:
        d = {}
        for i in range(len(files)):
            # drop _seq.rec suffix from path and take only filename
            _, filename = os.path.split(files[i][:-8])
            d[filename] = {'classes': classes[i].tolist(), 'Probabilities': probabilities[i].tolist()}
        json.dump(d, file)


def test_accuracy(model, folder_data, folder_results, print_summary, batch_size):
    """
    test already trained model on preprocessed data
    """
    y_test = np.array([], dtype=np.int32)
    files = list()
    folder_data_test = os.path.join(folder_data, 'seq_test')
    print('Getting file names for', model.num_classes, 'classes from folders:\n', folder_data_test)
    for i in range(1, model.num_classes + 1):
        folder = os.path.join(folder_data_test, str(i))
        assert os.path.exists(folder), "Folder: " + folder + ' does not exist'
        print('\ttest      : Read file names from folder ', folder)
        seq_files = [os.path.join(folder, f) for f in os.listdir(folder) if f[-4:] == '.rec']
        files += seq_files
        y_test = np.concatenate([y_test, np.array([int(i)] * len(seq_files), dtype=np.int32)])

    unk_index = model.unk_index
    X_seq_test, maxlen = encode_srcs(files, 'testing', unk_index)
    print('Max. sequence length overall:', maxlen)
    maxlen = model.maxlen if model.maxlen else maxlen
    print('Padding sequences to length', maxlen)
    X_seq_test = pad_src(X_seq_test, maxlen, unk_index)

    model_name = model.name
    predictions_path = model.predictions_path

    if fs.exists(predictions_path):
        print("\tFound predictions in", predictions_path, ", skipping...")
        with open(predictions_path, 'rb') as infile:
            classes = np.array(pickle.load(infile))
    else:
        model.restore()
        if print_summary:
            model.model.summary()

        print('\n--- Testing model...')
        classes, probabilities = model.predict(X_seq_test, batch_size)
        del probabilities
        fs.mkdir(fs.dirname(predictions_path))
        print('\tstore predictions to', predictions_path)
        with open(predictions_path, 'wb') as outfile:
            pickle.dump(classes.tolist(), outfile)

    print_stat(y_test, classes, model_name, folder_results)


def main():
    folder_results = FLAGS.folder_results
    folder_data = FLAGS.input_data
    print_summary = FLAGS.print_summary
    batch_size = FLAGS.batch_size

    embeddings = task_utils.get_embeddings()
    embedding_matrix_normalized = tf.nn.l2_normalize(embeddings, axis=1)

    with open(os.path.join(FLAGS.vocabulary_dir, 'dic_pickle'), 'rb') as f:
        dictionary = pickle.load(f)
    unk_index = dictionary[rgx.unknown_token]
    del dictionary

    model = NCC_classifyapp(seed=FLAGS.seed, maxlen=FLAGS.maxlen, num_classes=FLAGS.num_classes,
                            dense_layer_size=FLAGS.dense_layer_size, embedding_matrix=embedding_matrix_normalized,
                            folder_results=folder_results, model_name=FLAGS.model_name, unk_index=unk_index,
                            topk=FLAGS.topk, save_every=FLAGS.save_every, ring_size=FLAGS.ring_size)
    mode = FLAGS.mode
    if mode == 'predict' and FLAGS.json_file:
        predict_labels(model, FLAGS.json_file)
    elif mode == 'train':
        train_model(model, folder_data, folder_results, print_summary, FLAGS.num_epochs,
                    batch_size, FLAGS.trsamples, FLAGS.vasamples, FLAGS.tesamples)
    else:
        test_accuracy(model, folder_data, folder_results, print_summary, batch_size)


if __name__ == '__main__':
    main()