ydin0771's picture
Upload main.py
94566d7
raw
history blame
21.8 kB
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)