Upload main.py

main.py

@@ -0,0 +1,653 @@
+from __future__ import division
+import warnings
+
+from extract_feature import build_model, run_image, get_img_feat
+
+# warnings.filterwarnings("ignore", category=FutureWarning)
+# warnings.filterwarnings("ignore", message="size changed")
+warnings.filterwarnings("ignore")
+
+import sys
+import os
+import time
+import math
+import random
+
+try:
+    import Queue as queue
+except ImportError:
+    import queue
+import threading
+import h5py
+import json
+import numpy as np
+import tensorflow as tf
+from termcolor import colored, cprint
+
+from config import config, loadDatasetConfig, parseArgs
+from preprocess import Preprocesser, bold, bcolored, writeline, writelist
+from model import MACnet
+from collections import defaultdict
+
+
+############################################# loggers #############################################
+
+# Writes log header to file
+def logInit():
+    with open(config.logFile(), "a+") as outFile:
+        writeline(outFile, config.expName)
+        headers = ["epoch", "trainAcc", "valAcc", "trainLoss", "valLoss"]
+        if config.evalTrain:
+            headers += ["evalTrainAcc", "evalTrainLoss"]
+        if config.extra:
+            if config.evalTrain:
+                headers += ["thAcc", "thLoss"]
+            headers += ["vhAcc", "vhLoss"]
+        headers += ["time", "lr"]
+
+        writelist(outFile, headers)
+        # lr assumed to be last
+
+
+# Writes log record to file
+def logRecord(epoch, epochTime, lr, trainRes, evalRes, extraEvalRes):
+    with open(config.logFile(), "a+") as outFile:
+        record = [epoch, trainRes["acc"], evalRes["val"]["acc"], trainRes["loss"], evalRes["val"]["loss"]]
+        if config.evalTrain:
+            record += [evalRes["evalTrain"]["acc"], evalRes["evalTrain"]["loss"]]
+        if config.extra:
+            if config.evalTrain:
+                record += [extraEvalRes["evalTrain"]["acc"], extraEvalRes["evalTrain"]["loss"]]
+            record += [extraEvalRes["val"]["acc"], extraEvalRes["val"]["loss"]]
+        record += [epochTime, lr]
+
+        writelist(outFile, record)
+
+
+# Gets last logged epoch and learning rate
+def lastLoggedEpoch():
+    with open(config.logFile(), "r") as inFile:
+        lastLine = list(inFile)[-1].split(",")
+    epoch = int(lastLine[0])
+    lr = float(lastLine[-1])
+    return epoch, lr
+
+
+################################## printing, output and analysis ##################################
+
+# Analysis by type
+analysisQuestionLims = [(0, 18), (19, float("inf"))]
+analysisProgramLims = [(0, 12), (13, float("inf"))]
+
+toArity = lambda instance: instance["programSeq"][-1].split("_", 1)[0]
+toType = lambda instance: instance["programSeq"][-1].split("_", 1)[1]
+
+
+def fieldLenIsInRange(field):
+    return lambda instance, group: \
+        (len(instance[field]) >= group[0] and
+         len(instance[field]) <= group[1])
+
+
+# Groups instances based on a key
+def grouperKey(toKey):
+    def grouper(instances):
+        res = defaultdict(list)
+        for instance in instances:
+            res[toKey(instance)].append(instance)
+        return res
+
+    return grouper
+
+
+# Groups instances according to their match to condition
+def grouperCond(groups, isIn):
+    def grouper(instances):
+        res = {}
+        for group in groups:
+            res[group] = (instance for instance in instances if isIn(instance, group))
+        return res
+
+    return grouper
+
+
+groupers = {
+    "questionLength": grouperCond(analysisQuestionLims, fieldLenIsInRange("questionSeq")),
+    "programLength": grouperCond(analysisProgramLims, fieldLenIsInRange("programSeq")),
+    "arity": grouperKey(toArity),
+    "type": grouperKey(toType)
+}
+
+
+# Computes average
+def avg(instances, field):
+    if len(instances) == 0:
+        return 0.0
+    return sum(instances[field]) / len(instances)
+
+
+# Prints analysis of questions loss and accuracy by their group
+def printAnalysis(res):
+    if config.analysisType != "":
+        print("Analysis by {type}".format(type=config.analysisType))
+        groups = groupers[config.analysisType](res["preds"])
+        for key in groups:
+            instances = groups[key]
+            avgLoss = avg(instances, "loss")
+            avgAcc = avg(instances, "acc")
+            num = len(instances)
+            print("Group {key}: Loss: {loss}, Acc: {acc}, Num: {num}".format(key, avgLoss, avgAcc, num))
+
+
+# Print results for a tier
+def printTierResults(tierName, res, color):
+    if res is None:
+        return
+
+    print("{tierName} Loss: {loss}, {tierName} accuracy: {acc}".format(tierName=tierName,
+                                                                       loss=bcolored(res["loss"], color),
+                                                                       acc=bcolored(res["acc"], color)))
+
+    printAnalysis(res)
+
+
+# Prints dataset results (for several tiers)
+def printDatasetResults(trainRes, evalRes):
+    printTierResults("Training", trainRes, "magenta")
+    printTierResults("Training EMA", evalRes["evalTrain"], "red")
+    printTierResults("Validation", evalRes["val"], "cyan")
+
+
+# Writes predictions for several tiers
+def writePreds(preprocessor, evalRes):
+    preprocessor.writePreds(evalRes, "_")
+
+
+############################################# session #############################################
+# Initializes TF session. Sets GPU memory configuration.
+def setSession():
+    sessionConfig = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
+    if config.allowGrowth:
+        sessionConfig.gpu_options.allow_growth = True
+    if config.maxMemory < 1.0:
+        sessionConfig.gpu_options.per_process_gpu_memory_fraction = config.maxMemory
+    return sessionConfig
+
+
+############################################## savers #############################################
+# Initializes savers (standard, optional exponential-moving-average and optional for subset of variables)
+def setSavers(model):
+    saver = tf.train.Saver(max_to_keep=config.weightsToKeep)
+
+    subsetSaver = None
+    if config.saveSubset:
+        isRelevant = lambda var: any(s in var.name for s in config.varSubset)
+        relevantVars = [var for var in tf.global_variables() if isRelevant(var)]
+        subsetSaver = tf.train.Saver(relevantVars, max_to_keep=config.weightsToKeep, allow_empty=True)
+
+    emaSaver = None
+    if config.useEMA:
+        emaSaver = tf.train.Saver(model.emaDict, max_to_keep=config.weightsToKeep)
+
+    return {
+        "saver": saver,
+        "subsetSaver": subsetSaver,
+        "emaSaver": emaSaver
+    }
+
+
+################################### restore / initialize weights ##################################
+# Restores weights of specified / last epoch if on restore mod.
+# Otherwise, initializes weights.
+def loadWeights(sess, saver, init):
+    if config.restoreEpoch > 0 or config.restore:
+        # restore last epoch only if restoreEpoch isn't set
+        if config.restoreEpoch == 0:
+            # restore last logged epoch
+            config.restoreEpoch, config.lr = lastLoggedEpoch()
+        print(bcolored("Restoring epoch {} and lr {}".format(config.restoreEpoch, config.lr), "cyan"))
+        print(bcolored("Restoring weights", "blue"))
+        print(config.weightsFile(config.restoreEpoch))
+        saver.restore(sess, config.weightsFile(config.restoreEpoch))
+        epoch = config.restoreEpoch
+    else:
+        print(bcolored("Initializing weights", "blue"))
+        sess.run(init)
+        logInit()
+        epoch = 0
+
+    return epoch
+
+
+###################################### training / evaluation ######################################
+# Chooses data to train on (main / extra) data.
+def chooseTrainingData(data):
+    trainingData = data["main"]["train"]
+    alterData = None
+
+    if config.extra:
+        if config.trainExtra:
+            if config.extraVal:
+                trainingData = data["extra"]["val"]
+            else:
+                trainingData = data["extra"]["train"]
+        if config.alterExtra:
+            alterData = data["extra"]["train"]
+
+    return trainingData, alterData
+
+
+#### evaluation
+# Runs evaluation on train / val / test datasets.
+def runEvaluation(sess, model, data, epoch, evalTrain=True, evalTest=False, getAtt=None):
+    if getAtt is None:
+        getAtt = config.getAtt
+    res = {"evalTrain": None, "val": None, "test": None}
+
+    if data is not None:
+        if evalTrain and config.evalTrain:
+            res["evalTrain"] = runEpoch(sess, model, data["evalTrain"], train=False, epoch=epoch, getAtt=getAtt)
+
+        res["val"] = runEpoch(sess, model, data["val"], train=False, epoch=epoch, getAtt=getAtt)
+
+        if evalTest or config.test:
+            res["test"] = runEpoch(sess, model, data["test"], train=False, epoch=epoch, getAtt=getAtt)
+
+    return res
+
+
+## training conditions (comparing current epoch result to prior ones)
+def improveEnough(curr, prior, lr):
+    prevRes = prior["prev"]["res"]
+    currRes = curr["res"]
+
+    if prevRes is None:
+        return True
+
+    prevTrainLoss = prevRes["train"]["loss"]
+    currTrainLoss = currRes["train"]["loss"]
+    lossDiff = prevTrainLoss - currTrainLoss
+
+    notImprove = ((lossDiff < 0.015 and prevTrainLoss < 0.5 and lr > 0.00002) or \
+                  (lossDiff < 0.008 and prevTrainLoss < 0.15 and lr > 0.00001) or \
+                  (lossDiff < 0.003 and prevTrainLoss < 0.10 and lr > 0.000005))
+    # (prevTrainLoss < 0.2 and config.lr > 0.000015)
+
+    return not notImprove
+
+
+def better(currRes, bestRes):
+    return currRes["val"]["acc"] > bestRes["val"]["acc"]
+
+
+############################################## data ###############################################
+#### instances and batching
+# Trims sequences based on their max length.
+def trim2DVectors(vectors, vectorsLengths):
+    maxLength = np.max(vectorsLengths)
+    return vectors[:, :maxLength]
+
+
+# Trims batch based on question length.
+def trimData(data):
+    data["questions"] = trim2DVectors(data["questions"], data["questionLengths"])
+    return data
+
+
+# Gets batch / bucket size.
+def getLength(data):
+    return len(data["instances"])
+
+
+# Selects the data entries that match the indices.
+def selectIndices(data, indices):
+    def select(field, indices):
+        if type(field) is np.ndarray:
+            return field[indices]
+        if type(field) is list:
+            return [field[i] for i in indices]
+        else:
+            return field
+
+    selected = {k: select(d, indices) for k, d in data.items()}
+    return selected
+
+
+# Batches data into a a list of batches of batchSize.
+# Shuffles the data by default.
+def getBatches(data, batchSize=None, shuffle=True):
+    batches = []
+
+    dataLen = getLength(data)
+    if batchSize is None or batchSize > dataLen:
+        batchSize = dataLen
+
+    indices = np.arange(dataLen)
+    if shuffle:
+        np.random.shuffle(indices)
+
+    for batchStart in range(0, dataLen, batchSize):
+        batchIndices = indices[batchStart: batchStart + batchSize]
+        # if len(batchIndices) == batchSize?
+        if len(batchIndices) >= config.gpusNum:
+            batch = selectIndices(data, batchIndices)
+            batches.append(batch)
+            # batchesIndices.append((data, batchIndices))
+
+    return batches
+
+
+#### image batches
+# Opens image files.
+def openImageFiles(images):
+    images["imagesFile"] = h5py.File(images["imagesFilename"], "r")
+    images["imagesIds"] = None
+    if config.dataset == "NLVR":
+        with open(images["imageIdsFilename"], "r") as imageIdsFile:
+            images["imagesIds"] = json.load(imageIdsFile)
+
+        # Closes image files.
+
+
+def closeImageFiles(images):
+    images["imagesFile"].close()
+
+
+# Loads an images from file for a given data batch.
+def loadImageBatch(images, batch):
+    imagesFile = images["imagesFile"]
+    id2idx = images["imagesIds"]
+    toIndex = lambda imageId: imageId
+    if id2idx is not None:
+        toIndex = lambda imageId: id2idx[imageId]
+    imageBatch = np.stack([imagesFile["features"][toIndex(imageId)] for imageId in batch["imageIds"]], axis=0)
+
+    return {"images": imageBatch, "imageIds": batch["imageIds"]}
+
+
+# Loads images for several num batches in the batches list from start index.
+def loadImageBatches(images, batches, start, num):
+    batches = batches[start: start + num]
+    return [loadImageBatch(images, batch) for batch in batches]
+
+
+#### data alternation
+# Alternates main training batches with extra data.
+def alternateData(batches, alterData, dataLen):
+    alterData = alterData["data"][0]  # data isn't bucketed for altered data
+
+    # computes number of repetitions
+    needed = math.ceil(len(batches) / config.alterNum)
+    print(bold("Extra batches needed: %d") % needed)
+    perData = math.ceil(getLength(alterData) / config.batchSize)
+    print(bold("Batches per extra data: %d") % perData)
+    repetitions = math.ceil(needed / perData)
+    print(bold("reps: %d") % repetitions)
+
+    # make alternate batches
+    alterBatches = []
+    for _ in range(repetitions):
+        repBatches = getBatches(alterData, batchSize=config.batchSize)
+        random.shuffle(repBatches)
+        alterBatches += repBatches
+    print(bold("Batches num: %d") + len(alterBatches))
+
+    # alternate data with extra data
+    curr = len(batches) - 1
+    for alterBatch in alterBatches:
+        if curr < 0:
+            # print(colored("too many" + str(curr) + " " + str(len(batches)),"red"))
+            break
+        batches.insert(curr, alterBatch)
+        dataLen += getLength(alterBatch)
+        curr -= config.alterNum
+
+    return batches, dataLen
+
+
+############################################ threading ############################################
+
+imagesQueue = queue.Queue(maxsize=20)  # config.tasksNum
+inQueue = queue.Queue(maxsize=1)
+outQueue = queue.Queue(maxsize=1)
+
+
+# Runs a worker thread(s) to load images while training .
+class StoppableThread(threading.Thread):
+    # Thread class with a stop() method. The thread itself has to check
+    # regularly for the stopped() condition.
+
+    def __init__(self, images, batches):  # i
+        super(StoppableThread, self).__init__()
+        # self.i = i
+        self.images = images
+        self.batches = batches
+        self._stop_event = threading.Event()
+
+    # def __init__(self, args):
+    #     super(StoppableThread, self).__init__(args = args)
+    #     self._stop_event = threading.Event()
+
+    # def __init__(self, target, args):
+    #     super(StoppableThread, self).__init__(target = target, args = args)
+    #     self._stop_event = threading.Event()
+
+    def stop(self):
+        self._stop_event.set()
+
+    def stopped(self):
+        return self._stop_event.is_set()
+
+    def run(self):
+        while not self.stopped():
+            try:
+                batchNum = inQueue.get(timeout=60)
+                nextItem = loadImageBatches(self.images, self.batches, batchNum, int(config.taskSize / 2))
+                outQueue.put(nextItem)
+                # inQueue.task_done()
+            except:
+                pass
+        # print("worker %d done", self.i)
+
+
+def loaderRun(images, batches):
+    batchNum = 0
+
+    # if config.workers == 2:
+    #     worker = StoppableThread(images, batches) # i,
+    #     worker.daemon = True
+    #     worker.start()
+
+    #     while batchNum < len(batches):
+    #         inQueue.put(batchNum + int(config.taskSize / 2))
+    #         nextItem1 = loadImageBatches(images, batches, batchNum, int(config.taskSize / 2))
+    #         nextItem2 = outQueue.get()
+
+    #         nextItem = nextItem1 + nextItem2
+    #         assert len(nextItem) == min(config.taskSize, len(batches) - batchNum)
+    #         batchNum += config.taskSize
+
+    #         imagesQueue.put(nextItem)
+
+    #     worker.stop()
+    # else:
+    while batchNum < len(batches):
+        nextItem = loadImageBatches(images, batches, batchNum, config.taskSize)
+        assert len(nextItem) == min(config.taskSize, len(batches) - batchNum)
+        batchNum += config.taskSize
+        imagesQueue.put(nextItem)
+
+    # print("manager loader done")
+
+
+########################################## stats tracking #########################################
+# Computes exponential moving average.
+def emaAvg(avg, value):
+    if avg is None:
+        return value
+    emaRate = 0.98
+    return avg * emaRate + value * (1 - emaRate)
+
+
+# Initializes training statistics.
+def initStats():
+    return {
+        "totalBatches": 0,
+        "totalData": 0,
+        "totalLoss": 0.0,
+        "totalCorrect": 0,
+        "loss": 0.0,
+        "acc": 0.0,
+        "emaLoss": None,
+        "emaAcc": None,
+    }
+
+
+# Updates statistics with training results of a batch
+def updateStats(stats, res, batch):
+    stats["totalBatches"] += 1
+    stats["totalData"] += getLength(batch)
+
+    stats["totalLoss"] += res["loss"]
+    stats["totalCorrect"] += res["correctNum"]
+
+    stats["loss"] = stats["totalLoss"] / stats["totalBatches"]
+    stats["acc"] = stats["totalCorrect"] / stats["totalData"]
+
+    stats["emaLoss"] = emaAvg(stats["emaLoss"], res["loss"])
+    stats["emaAcc"] = emaAvg(stats["emaAcc"], res["acc"])
+
+    return stats
+
+
+# auto-encoder ae = {:2.4f} autoEncLoss,
+# Translates training statistics into a string to print
+def statsToStr(stats, res, epoch, batchNum, dataLen, startTime):
+    formatStr = "\reb {epoch},{batchNum} ({dataProcessed} / {dataLen:5d}), " + \
+                "t = {time} ({loadTime:2.2f}+{trainTime:2.2f}), " + \
+                "lr {lr}, l = {loss}, a = {acc}, avL = {avgLoss}, " + \
+                "avA = {avgAcc}, g = {gradNorm:2.4f}, " + \
+                "emL = {emaLoss:2.4f}, emA = {emaAcc:2.4f}; " + \
+                "{expname}"  # {machine}/{gpu}"
+
+    s_epoch = bcolored("{:2d}".format(epoch), "green")
+    s_batchNum = "{:3d}".format(batchNum)
+    s_dataProcessed = bcolored("{:5d}".format(stats["totalData"]), "green")
+    s_dataLen = dataLen
+    s_time = bcolored("{:2.2f}".format(time.time() - startTime), "green")
+    s_loadTime = res["readTime"]
+    s_trainTime = res["trainTime"]
+    s_lr = bold(config.lr)
+    s_loss = bcolored("{:2.4f}".format(res["loss"]), "blue")
+    s_acc = bcolored("{:2.4f}".format(res["acc"]), "blue")
+    s_avgLoss = bcolored("{:2.4f}".format(stats["loss"]), "blue")
+    s_avgAcc = bcolored("{:2.4f}".format(stats["acc"]), "red")
+    s_gradNorm = res["gradNorm"]
+    s_emaLoss = stats["emaLoss"]
+    s_emaAcc = stats["emaAcc"]
+    s_expname = config.expName
+    # s_machine = bcolored(config.dataPath[9:11],"green")
+    # s_gpu = bcolored(config.gpus,"green")
+
+    return formatStr.format(epoch=s_epoch, batchNum=s_batchNum, dataProcessed=s_dataProcessed,
+                            dataLen=s_dataLen, time=s_time, loadTime=s_loadTime,
+                            trainTime=s_trainTime, lr=s_lr, loss=s_loss, acc=s_acc,
+                            avgLoss=s_avgLoss, avgAcc=s_avgAcc, gradNorm=s_gradNorm,
+                            emaLoss=s_emaLoss, emaAcc=s_emaAcc, expname=s_expname)
+    # machine = s_machine, gpu = s_gpu)
+
+
+# collectRuntimeStats, writer = None,
+'''
+Runs an epoch with model and session over the data.
+1. Batches the data and optionally mix it with the extra alterData.
+2. Start worker threads to load images in parallel to training.
+3. Runs model for each batch, and gets results (e.g. loss,  accuracy).
+4. Updates and prints statistics based on batch results.
+5. Once in a while (every config.saveEvery), save weights. 
+
+Args:
+    sess: TF session to run with.
+
+    model: model to process data. Has runBatch method that process a given batch.
+    (See model.py for further details).
+
+    data: data to use for training/evaluation.
+
+    epoch: epoch number.
+
+    saver: TF saver to save weights
+
+    calle: a method to call every number of iterations (config.calleEvery)
+
+    alterData: extra data to mix with main data while training.
+
+    getAtt: True to return model attentions.  
+'''
+
+
+def main(question, image):
+    with open(config.configFile(), "a+") as outFile:
+        json.dump(vars(config), outFile)
+
+    # set gpus
+    if config.gpus != "":
+        config.gpusNum = len(config.gpus.split(","))
+        os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus
+
+    tf.logging.set_verbosity(tf.logging.ERROR)
+
+    # process data
+    print(bold("Preprocess data..."))
+    start = time.time()
+    preprocessor = Preprocesser()
+    cnn_model = build_model()
+    imageData = get_img_feat(cnn_model, image)
+    qData, embeddings, answerDict = preprocessor.preprocessData(question)
+    data = {'data': qData, 'image': imageData}
+    print("took {} seconds".format(bcolored("{:.2f}".format(time.time() - start), "blue")))
+
+    # build model
+    print(bold("Building model..."))
+    start = time.time()
+    model = MACnet(embeddings, answerDict)
+    print("took {} seconds".format(bcolored("{:.2f}".format(time.time() - start), "blue")))
+
+    # initializer
+    init = tf.global_variables_initializer()
+
+    # savers
+    savers = setSavers(model)
+    saver, emaSaver = savers["saver"], savers["emaSaver"]
+
+    # sessionConfig
+    sessionConfig = setSession()
+
+    with tf.Session(config=sessionConfig) as sess:
+
+        # ensure no more ops are added after model is built
+        sess.graph.finalize()
+
+        # restore / initialize weights, initialize epoch variable
+        epoch = loadWeights(sess, saver, init)
+        print(epoch)
+        start = time.time()
+        if epoch > 0:
+            if config.useEMA:
+                emaSaver.restore(sess, config.weightsFile(epoch))
+            else:
+                saver.restore(sess, config.weightsFile(epoch))
+
+        evalRes = model.runBatch(sess, data['data'], data['image'], False)
+
+        print("took {:.2f} seconds".format(time.time() - start))
+
+        print(evalRes)
+
+
+if __name__ == '__main__':
+    parseArgs()
+    loadDatasetConfig[config.dataset]()
+    question = 'How many text objects are located at the bottom side of table?'
+    imagePath = './mac-layoutLM-sample/PDF_val_64.png'
+    main(question, imagePath)