def __init__(self, id, args, worker_address, sink_address, device_id):
super().__init__()
self.model_dir = args.model_dir
self.config_fp = os.path.join(self.model_dir, 'bert_config.json')
self.checkpoint_fp = os.path.join(self.model_dir, 'bert_model.ckpt')
self.vocab_fp = os.path.join(args.model_dir, 'vocab.txt')
self.tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_fp)
self.max_seq_len = args.max_seq_len
self.worker_id = id
self.daemon = True
self.model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(self.config_fp),
init_checkpoint=self.checkpoint_fp,
pooling_strategy=args.pooling_strategy,
pooling_layer=args.pooling_layer,
use_xla=args.xla
)
os.environ['CUDA_VISIBLE_DEVICES'] = str(device_id)
config = tf.ConfigProto(device_count={'GPU': 0 if device_id < 0 else 1})
# session-wise XLA doesn't seem to work on tf 1.10
# if args.xla:
# config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = args.gpu_memory_fraction
self.estimator = Estimator(self.model_fn, config=RunConfig(session_config=config))
self.exit_flag = multiprocessing.Event()
self.logger = set_logger(colored('WORKER-%d' % self.worker_id, 'yellow'))
self.worker_address = worker_address
self.sink_address = sink_address
self.prefetch_factor = 10
def generate_results(model_dir, results_path):
os.makedirs(os.path.dirname(results_path), exist_ok=True)
vocab = np.load('output/processed-annotations/vocab.npy')
run_config = RunConfig(model_dir=model_dir)
hparams = get_hparams(model_dir=model_dir, create=False)
print(hparams)
estimator = Estimator(model_fn=model_fn, config=run_config, params=hparams)
val_path = tf.flags.FLAGS.batch_path
splits = tf.flags.FLAGS.batch_splits
batch_size = tf.flags.FLAGS.batch_size
hook = FeedFnHook(path_fmt=val_path,
splits=splits,
batch_size=batch_size,
predict=True,
single_pass=True)
results = []
it = tqdm(desc='Generating results')
for prediction in estimator.predict(input_fn=predict_input_fn,
hooks=[hook]):
caption = calc_caption(prediction=prediction, vocab=vocab)
results.append({
'image_id': np.asscalar(prediction['image_ids']),
'caption': caption
})
it.update(1)
with open(results_path, 'w') as f:
json.dump(results, f)
def __init__(self, id, args, worker_address, sink_address):
super().__init__()
self.model_dir = args.model_dir
self.config_fp = os.path.join(self.model_dir, 'bert_config.json')
self.checkpoint_fp = os.path.join(self.model_dir, 'bert_model.ckpt')
self.vocab_fp = os.path.join(args.model_dir, 'vocab.txt')
self.tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_fp)
self.max_seq_len = args.max_seq_len
self.worker_id = id
self.daemon = True
self.model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(self.config_fp),
init_checkpoint=self.checkpoint_fp,
pooling_strategy=args.pooling_strategy,
pooling_layer=args.pooling_layer)
os.environ['CUDA_VISIBLE_DEVICES'] = str(self.worker_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = args.gpu_memory_fraction
self.estimator = Estimator(self.model_fn,
config=RunConfig(session_config=config))
self.exit_flag = multiprocessing.Event()
self.logger = set_logger('WORKER-%d' % self.worker_id)
self.worker_address = worker_address
self.sink_address = sink_address
def __init__(self, sen2vec_conf_file):
tf.logging.set_verbosity(tf.logging.WARN)
##Sentence2Vec配置
self.sen2vec_conf = Sentence2VecConfig.from_json_file(
sen2vec_conf_file)
##bert模型+分词器
self.bert_config = modeling.BertConfig.from_json_file(
self.sen2vec_conf.bert_config_file)
self.tokenizer = tokenization.FullTokenizer(
vocab_file=self.sen2vec_conf.vocab_file,
do_lower_case=self.sen2vec_conf.do_lower_case)
self.model_fn = self.model_fn_builder(
bert_config=self.bert_config,
init_checkpoint=self.sen2vec_conf.init_checkpoint,
use_one_hot_embeddings=self.sen2vec_conf.use_one_hot_embeddings,
pooling_layer=self.sen2vec_conf.feature_pooling_layer)
# Remove TPU Estimator.
#estimator = tf.contrib.tpu.TPUEstimator(
# use_tpu=self.sen2vec_conf.use_tpu,
# model_fn=model_fn,
# config=run_config,
# predict_batch_size=self.sen2vec_conf.batch_size)
self.params = {}
self.params["batch_size"] = self.sen2vec_conf.batch_size
self.estimator = Estimator(self.model_fn, params=self.params)
self.seq_length = self.sen2vec_conf.max_seq_length
def get_estimator(self, tf):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
from tensorflow.python.estimator.model_fn import EstimatorSpec
def model_fn(features, labels, mode, params):
with tf.gfile.GFile(self.graph_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
input_names = ['input_ids', 'input_mask', 'input_type_ids']
output = tf.import_graph_def(graph_def,
input_map={k + ':0': features[k] for k in input_names},
return_elements=['final_encodes:0'])
return EstimatorSpec(mode=mode, predictions={
'client_id': features['client_id'],
'encodes': output[0]
})
def ner_model_fn(features, labels, mode, params):
"""
命名实体识别模型的model_fn
:param features:
:param labels:
:param mode:
:param params:
:return:
"""
with tf.gfile.GFile(self.graph_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
input_ids = features["input_ids"]
input_mask = features["input_mask"]
input_map = {"input_ids": input_ids, "input_mask": input_mask}
pred_ids = tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['pred_ids:0'])
return EstimatorSpec(mode=mode, predictions={
'client_id': features['client_id'],
'encodes': pred_ids[0]
})
# 0 表示只使用CPU 1 表示使用GPU
config = tf.ConfigProto(device_count={'GPU': 0 if self.device_id < 0 else 1})
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = self.gpu_memory_fraction
config.log_device_placement = False
# session-wise XLA doesn't seem to work on tf 1.10
# if args.xla:
# config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
if self.mode == 'NER':
return Estimator(model_fn=ner_model_fn, config=RunConfig(session_config=config))
elif self.mode == 'BERT':
return Estimator(model_fn=model_fn, config=RunConfig(session_config=config))
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
# Get corpus info
FLAGS.n_token = data_utils.VOCAB_SIZE
tf.logging.info('n_token {}'.format(FLAGS.n_token))
if not tf.gfile.Exists(FLAGS.model_dir):
tf.gfile.MakeDirs(FLAGS.model_dir)
bsz_per_core = per_device_batch_size(FLAGS.train_batch_size,
FLAGS.num_gpu_cores)
tf.logging.info('size of batch {}'.format(bsz_per_core))
train_input_fn, train_record_info_dict = get_input_fn(
'train', bsz_per_core)
tf.logging.info('num of batches {}'.format(
train_record_info_dict['num_batch']))
train_cache_fn = get_cache_fn(FLAGS.mem_len, bsz_per_core)
tf.logging.info(train_cache_fn)
tf.logging.info('Use normal RunConfig')
tf.logging.info(FLAGS.num_gpu_cores)
dist_strategy = tf.contrib.distribute.MirroredStrategy(
num_gpus=FLAGS.num_gpu_cores,
auto_shard_dataset=True,
cross_device_ops=AllReduceCrossDeviceOps(
'nccl', num_packs=FLAGS.num_gpu_cores),
# cross_device_ops=AllReduceCrossDeviceOps('hierarchical_copy'),
)
log_every_n_steps = 10
run_config = RunConfig(
train_distribute=dist_strategy,
eval_distribute=dist_strategy,
log_step_count_steps=log_every_n_steps,
model_dir=FLAGS.model_dir,
save_checkpoints_steps=FLAGS.save_steps,
save_summary_steps=None,
)
model_fn = get_model_fn()
tf.logging.info('Use normal Estimator')
estimator = Estimator(
model_fn=model_fn,
params={
'batch_size': bsz_per_core,
'cache': None
},
config=run_config,
)
tf.logging.info('***** Running training *****')
tf.logging.info(' Batch size = %d', FLAGS.train_batch_size)
estimator.train(input_fn=train_input_fn, max_steps=FLAGS.train_steps)
def main():
sequence_schema_path = f'{input_path}/train/sequence_schema'
context_schema_path = f'{input_path}/train/context_schema'
context_schema, sequence_schema = read_schemata(context_schema_path, sequence_schema_path)
tf_ctx_schema, tf_seq_schema = build_schema(context_schema, sequence_schema)
train_parts = glob.glob(input_path + '/train' + '/part-*')
validation_parts = glob.glob(input_path + '/test' + '/part-*')
run_config = RunConfig(log_step_count_steps=10,
save_checkpoints_steps=100,
save_summary_steps=200,
keep_checkpoint_max=32)
shared_input_fn = partial(input_fn, params, tf_seq_schema, tf_ctx_schema)
train_input_fn = partial(shared_input_fn, train_parts)
validation_input_fn = partial(shared_input_fn, validation_parts)
train_spec = TrainSpec(train_input_fn, max_steps=1000000)
eval_spec = EvalSpec(validation_input_fn, steps=200, name='validation', start_delay_secs=30, throttle_secs=1)
estimator = Estimator(model_fn=model.model_fn,
model_dir=model_dir,
params=params,
config=run_config)
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=logging.INFO)
logging.getLogger('tensorflow').propagate = False
train_and_evaluate(estimator=estimator,
train_spec=train_spec,
eval_spec=eval_spec)
prediction = list(estimator.predict(input_fn=partial(predict_input_fn, {'epochs': 1, 'batch_size': 10}, grid)))
scores = [p.tolist() for p in prediction]
pairwise_prob = pairwise_probability(scores)
zero = pairwise_prob[0]
A_zero = build_diags(zero)
print(optimize(A_zero).x)
def process(question, contexts):
# TODO Replace all abbreviation code
bert_config = modeling.BertConfig.from_json_file(
os.path.join(modelDir, 'bert_config.json')) # Loading bert config
tokenizer = tokenization.FullTokenizer(
vocab_file=os.path.join(modelDir, 'vocab.txt'),
do_lower_case=False) # Loading tokenizer
candidates = read_QA(question, contexts)
eval_features = convert_candidates_to_features(candidates=candidates,
tokenizer=tokenizer,
max_seq_length=512,
doc_stride=256,
max_query_length=128)
model_fn = model_fn_builder(bert_config=bert_config,
init_checkpoint=os.path.join(
modelDir, 'bert_model.ckpt'),
use_one_hot_embeddings=False)
run_config = RunConfig(model_dir=modelDir, save_checkpoints_steps=1000)
estimator = Estimator(model_fn=model_fn,
config=run_config,
params={'batch_size': 14})
predict_input_fn = input_fn_builder(features=eval_features,
seq_length=512,
drop_remainder=True)
all_results = []
counter = 0
RawResult = collections.namedtuple(
"RawResult", ["unique_id", "start_logits", "end_logits"])
for result in estimator.predict(predict_input_fn,
yield_single_examples=True):
unique_id = int(result["unique_ids"])
start_logits = [float(x) for x in result["start_logits"].flat]
end_logits = [float(x) for x in result["end_logits"].flat]
all_results.append(
RawResult(unique_id=unique_id,
start_logits=start_logits,
end_logits=end_logits))
counter += 1
if len(eval_features) == counter: break
all_nbest_json = write_QA(candidates, eval_features, all_results, 2, 128,
False)
return all_nbest_json
def create_estimator_and_tokenizer(self):
config_fp = os.path.join(self.model_dir, 'bert_config.json')
checkpoint_fp = os.path.join(self.model_dir, 'bert_model.ckpt')
vocab_fp = os.path.join(self.model_dir, 'vocab.txt')
tokenizer = tokenization.FullTokenizer(vocab_file=vocab_fp)
model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(config_fp),
init_checkpoint=checkpoint_fp,
pooling_strategy=PoolingStrategy.NONE,
pooling_layer=[-2]
)
config = tf.ConfigProto(
device_count={
'CPU': cpu_count()
},
inter_op_parallelism_threads=0,
intra_op_parallelism_threads=0,
gpu_options={
'allow_growth': True
}
)
estimator = Estimator(model_fn, config=RunConfig(
session_config=config), model_dir=None)
return estimator, tokenizer
def get_estimator(self):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
bert_config = modeling.BertConfig.from_json_file(args.config_name)
label_list = self.processor.get_labels()
train_examples = self.processor.get_train_examples(args.data_dir)
num_train_steps = int(
len(train_examples) / self.batch_size * args.num_train_epochs)
num_warmup_steps = int(num_train_steps * 0.1)
if self.mode == tf.estimator.ModeKeys.TRAIN:
init_checkpoint = args.ckpt_name
else:
init_checkpoint = args.output_dir
model_fn = self.model_fn_builder(bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=init_checkpoint,
learning_rate=args.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_one_hot_embeddings=False)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = args.gpu_memory_fraction
config.log_device_placement = False
return Estimator(model_fn=model_fn,
config=RunConfig(session_config=config),
model_dir=args.output_dir,
params={'batch_size': self.batch_size})
def get_estimator(self, tf):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
from tensorflow.python.estimator.model_fn import EstimatorSpec
def model_fn(features, labels, mode, params):
with tf.gfile.GFile(self.graph_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
input_ids = features["input_ids"]
input_mask = features["input_mask"]
input_map = {"input_ids": input_ids, "input_mask": input_mask}
pred_theme, embeddings = tf.import_graph_def(
graph_def, name='', input_map=input_map,
return_elements=['theme/logits/ArgMax:0', 'bert/transformer/Reshape_2:0'])
return EstimatorSpec(mode=mode, predictions={
'client_id': features['client_id'],
'pred_themes': pred_theme,
'embeddings': embeddings
})
# 0 表示只使用CPU 1 表示使用GPU
config = tf.ConfigProto(device_count={'GPU': 0 if self.device_id < 0 else 1})
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = self.gpu_memory_fraction
config.log_device_placement = False
return Estimator(model_fn=model_fn, config=RunConfig(session_config=config))
def get_estimator(self):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
bert_config = modeling.BertConfig.from_json_file(args.config_name)
label_list = self.processor.get_labels()
train_examples = self.processor.get_train_examples(args.data_dir)
num_train_steps = int(
len(train_examples) / self.batch_size * args.num_train_epochs)
#学习率预热系数
num_warmup_steps = int(num_train_steps * 0.1)
#模型框架初始值，如果是微调训练模型初始值使用bert中的initpoint如果是预测模型的参数使用output中的
if self.mode == tf.estimator.ModeKeys.TRAIN:
init_checkpoint = args.ckpt_name
else:
init_checkpoint = args.output_dir
#创建model_fn对象,该对象中定义了整个模型结构框架
model_fn = self.model_fn_builder(bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=init_checkpoint,
learning_rate=args.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_one_hot_embeddings=False)
#config = tf.ConfigProto()
#config.gpu_options.allow_growth = True
#config.gpu_options.per_process_gpu_memory_fraction = args.gpu_memory_fraction
strategy = tf.contrib.distribute.MirroredStrategy(
devices=["/device:GPU:0", "/device:GPU:1"])
config_dis = tf.estimator.RunConfig(train_distribute=strategy)
return Estimator(model_fn=model_fn,
config=config_dis,
model_dir=args.output_dir,
params={'batch_size': self.batch_size})
def get_estimator(bert_config_file, init_checkpoint, max_seq_len, select_layers, graph_file = '../bert/tmp/graph'):
#from tensorflow.python.estimator.estimator import Estimator
#from tensorflow.python.estimator.run_config import RunConfig
#from tensorflow.python.estimator.model_fn import EstimatorSpec
if os.path.exists(graph_file):
graph_path =graph_file
else:
graph_path = create_graph(graph_file, bert_config_file, init_checkpoint, max_seq_len, select_layers)
def model_fn(features, labels, mode, params):
with tf.gfile.GFile(graph_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
input_names = ['input_ids', 'input_mask', 'input_type_ids']
encoder_layer = tf.import_graph_def(graph_def,
input_map={k + ':0': features[k] for k in input_names},
return_elements=['final_encodes:0'])
predictions = {
# 'client_id': client_id,
'encodes': encoder_layer[0]
}
return EstimatorSpec(mode=mode, predictions=predictions)
config = tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)
config.gpu_options.allow_growth = True
#config.gpu_options.per_process_gpu_memory_fraction = self.gpu_memory_fraction
#config.log_device_placement = False
#config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
return Estimator(model_fn=model_fn, config=RunConfig(session_config=config),
params={'batch_size': 32}, model_dir='../bert/tmp')
def get_estimator(self):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
from tensorflow.python.estimator.model_fn import EstimatorSpec
def model_fn(features, labels, mode, params):
with tf.gfile.GFile(self.graph_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
input_names = ['input_ids', 'input_mask', 'input_type_ids']
output = tf.import_graph_def(
graph_def,
input_map={k + ':0': features[k]
for k in input_names},
return_elements=['final_encodes:0'])
return EstimatorSpec(mode=mode, predictions={'encodes': output[0]})
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = self.gpu_memory_fraction
config.log_device_placement = False
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
return Estimator(model_fn=model_fn,
config=RunConfig(session_config=config),
params={'batch_size': self.batch_size})
def experiment_fn(run_config, hparams):
if hparams.dataset == "COCO":
dataset = COCO("data/coco")
elif hparams.dataset == "challenger.ai":
dataset = ChallengerAI("data/challenger.ai")
else:
raise Exception("Unknown Dataset Name: '%s'." % hparams.dataset)
hparams.add_hparam("vocab_size", len(dataset.word_to_idx))
estimator = Estimator(model_fn=model_fn, params=hparams, config=run_config)
train_init_hook = IteratorInitializerHook("train")
val_init_hook = IteratorInitializerHook("val")
experiment = Experiment(
estimator=estimator,
train_input_fn=dataset.get_input_fn("train", is_distort=True),
eval_input_fn=dataset.get_input_fn("val"),
train_steps=hparams.train_steps,
eval_steps=hparams.eval_steps,
train_steps_per_iteration=hparams.steps_per_eval,
eval_hooks=[val_init_hook],
)
experiment.extend_train_hooks([train_init_hook])
return experiment
def train():
model_dir = tf.flags.FLAGS.model_dir
os.makedirs(model_dir, exist_ok=True)
print("model_dir={}".format(model_dir))
run_config = RunConfig(
model_dir=model_dir,
save_checkpoints_steps=tf.flags.FLAGS.save_checkpoints_steps)
hparams = get_hparams(model_dir, validate=True)
vocabs = read_vocablists(path=tf.flags.FLAGS.data_dir)
train_input_fn, eval_input_fn, test_input_fn = make_input_fns(
tf.flags.FLAGS.data_dir, batch_size=tf.flags.FLAGS.batch_size)
# Model
model_fn = make_model_fn(hparams=hparams,
run_config=run_config,
vocabs=vocabs)
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn,
max_steps=tf.flags.FLAGS.max_steps)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn,
steps=tf.flags.FLAGS.eval_steps,
throttle_secs=0)
estimator = Estimator(model_fn=model_fn, config=run_config, params=hparams)
train_and_evaluate(eval_spec=eval_spec,
train_spec=train_spec,
estimator=estimator)
def __init__(self, id, args):
super().__init__()
self.model_dir = args.model_dir
self.config_fp = os.path.join(self.model_dir, 'bert_config.json')
self.checkpoint_fp = os.path.join(self.model_dir, 'bert_model.ckpt')
self.vocab_fp = os.path.join(args.model_dir, 'vocab.txt')
self.tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_fp)
self.max_len = args.max_len
self.worker_id = id
self.daemon = True
self.model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(self.config_fp),
init_checkpoint=self.checkpoint_fp)
os.environ['CUDA_VISIBLE_DEVICES'] = str(self.worker_id)
self.estimator = Estimator(self.model_fn)
self.result = []
def train():
model_dir = FLAGS.model_dir
os.makedirs(model_dir, exist_ok=True)
hparams = get_hparams(
model_dir,
DEFAULT_HPARAMS,
hparams_file=FLAGS.config,
hparams_str=FLAGS.hparams,
validate=True
)
run_config = RunConfig(
model_dir=model_dir,
save_checkpoints_steps=FLAGS.save_checkpoints_steps,
save_summary_steps=FLAGS.save_summary_steps)
train_input_fn, eval_input_fn = make_input_fns()
# Model
estimator = Estimator(
model_fn=model_fn,
config=run_config,
params=hparams)
train_spec = tf.estimator.TrainSpec(
input_fn=train_input_fn,
max_steps=FLAGS.max_steps)
eval_spec = tf.estimator.EvalSpec(
input_fn=eval_input_fn,
steps=FLAGS.eval_steps,
throttle_secs=0)
train_and_evaluate(
eval_spec=eval_spec,
train_spec=train_spec,
estimator=estimator
)
def get_estimator(self, tf):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
from tensorflow.python.estimator.model_fn import EstimatorSpec
def model_fn(features, labels, mode, params):
with tf.io.gfile.GFile(self.graph_path, 'rb') as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
input_names = ['input_ids', 'input_mask', 'input_type_ids']
output = tf.import_graph_def(graph_def,
input_map={k + ':0': features[k] for k in input_names},
return_elements=['final_encodes:0'])
return EstimatorSpec(mode=mode, predictions={
'client_id': features['client_id'],
'encodes': output[0]
})
config = tf.compat.v1.ConfigProto(device_count={'GPU': 0 if self.device_id < 0 else 1})
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = self.gpu_memory_fraction
config.log_device_placement = False
# session-wise XLA doesn't seem to work on tf 1.10
# if args.xla:
# config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
return Estimator(model_fn=model_fn, config=RunConfig(session_config=config))
def experiment_fn(run_config, hparams):
eval_initializer_hook = IteratorInitializerHook()
estimator = Estimator(model_fn=model_fn, params=hparams, config=run_config)
train_dataset = CornellMovieDataset(
os.path.join(config.PROCESSED_PATH, "train_ids.enc"),
os.path.join(config.PROCESSED_PATH, "train_ids.dec"),
os.path.join(config.PROCESSED_PATH, "vocab.enc"),
os.path.join(config.PROCESSED_PATH, "vocab.dec"),
)
test_dataset = CornellMovieDataset(
os.path.join(config.PROCESSED_PATH, "test_ids.enc"),
os.path.join(config.PROCESSED_PATH, "test_ids.dec"),
os.path.join(config.PROCESSED_PATH, "vocab.enc"),
os.path.join(config.PROCESSED_PATH, "vocab.dec"),
)
experiment = Experiment(estimator=estimator,
train_input_fn=train_dataset.input_fn,
eval_input_fn=test_dataset.input_fn,
train_steps=hparams.train_steps,
eval_hooks=[eval_initializer_hook],
eval_steps=None,
train_steps_per_iteration=hparams.steps_per_eval)
return experiment
def get_estimator(self):
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig
bert_config = modeling.BertConfig.from_json_file(cf.bert_config_file)
train_examples = self.processor.get_train_examples(cf.data_dir)
label_list = self.processor.get_labels() # 这里需要这样写，如果用self.get_label_list()获取列表，在还没有生成label_list.pkl文件的时候会报错
# label_list = self.get_label_list()
num_train_steps = int(len(train_examples) / self.batch_size * cf.num_train_epochs)
num_warmup_steps = int(num_train_steps * 0.1)
if self.mode == tf.estimator.ModeKeys.TRAIN:
init_checkpoint = cf.init_checkpoint
else:
init_checkpoint = cf.output_dir # 预测模式下加载
model_fn = self.model_fn_builder(
bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=init_checkpoint,
learning_rate=cf.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_one_hot_embeddings=False)
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = cf.gpu_memory_fraction
config.log_device_placement = False
return Estimator(model_fn=model_fn, config=RunConfig(session_config=config), model_dir=cf.output_dir, params={'batch_size': self.batch_size})
def experiment_fn(run_config, hparams):
train_input_fn, eval_input_fn = make_input_fns()
estimator = Estimator(model_fn=model_fn, config=run_config, params=hparams)
experiment = Experiment(estimator=estimator,
train_input_fn=train_input_fn,
eval_input_fn=eval_input_fn)
return experiment
def get_estimator():
config = tf.ConfigProto(device_count={'GPU': 0})
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.log_device_placement = False
return Estimator(model_fn=model_fn,
config=RunConfig(session_config=config))
def main(_):
tf.logging.set_verbosity(tf.logging.WARN)
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
examples = read_examples(FLAGS.input_file)
features = convert_examples_to_features(examples=examples,
seq_length=FLAGS.max_seq_length,
tokenizer=tokenizer)
unique_id_to_feature = {}
for feature in features:
unique_id_to_feature[feature.unique_id] = feature
model_fn = model_fn_builder(
bert_config=bert_config,
init_checkpoint=FLAGS.init_checkpoint,
use_one_hot_embeddings=FLAGS.use_one_hot_embeddings,
pooling_layer=FLAGS.feature_pooling_layer)
# Remove TPU Estimator.
#estimator = tf.contrib.tpu.TPUEstimator(
# use_tpu=FLAGS.use_tpu,
# model_fn=model_fn,
# config=run_config,
# predict_batch_size=FLAGS.batch_size)
estimator = Estimator(model_fn, params=pack_params())
input_fn = input_fn_builder(features=features,
seq_length=FLAGS.max_seq_length)
with codecs.getwriter("utf-8")(tf.gfile.Open(FLAGS.output_file,
"w")) as writer:
for result in estimator.predict(input_fn, yield_single_examples=True):
unique_id = int(result["unique_id"])
encodes = [round(float(x), 6) for x in result["encodes"].flat]
feature = unique_id_to_feature[unique_id]
text = feature.tokens[1:-1]
#tf.logging.info("str: %s, encodes[%d]=%s" % ("".join(text), len(encodes), ",".join([str(x) for x in encodes])))
writer.write("\t".join([
str(unique_id), "".join(text), ",".join(
[str(x) for x in encodes])
]) + "\n")
def experiment_fn(run_config, hparams):
estimator = Estimator(model_fn=config.model_fn,
config=run_config,
params=hparams)
experiment = Experiment(estimator=estimator,
train_input_fn=config.input_fns['train'],
eval_input_fn=config.input_fns['valid'])
return experiment
def main():
parsed_args = get_parser().parse_args()
with open(os.path.join("data/challenger.ai", 'word_to_idx.pkl'), 'rb') as f:
word_to_idx = pickle.load(f)
hparams = HParams(vocab_size=len(word_to_idx),
batch_size=parsed_args.batch_size,
selector=parsed_args.selector,
dropout=parsed_args.dropout,
ctx2out=parsed_args.ctx2out,
prev2out=parsed_args.prev2out,
hard_attention=parsed_args.hard_attention,
bin_size=14)
run_config = RunConfig(model_dir=parsed_args.model_dir)
estimator = Estimator(
model_fn=model_fn_inner,
params=hparams,
config=run_config)
dataset = ChallengerAI("data/challenger.ai")
input_fn = dataset.get_tfrecords_test_input_fn(bin_size=hparams.bin_size)
val_init_hook = IteratorInitializerHook("infer")
idx_to_word = {v: k for k, v in word_to_idx.items()}
del word_to_idx
results = estimator.predict(input_fn, hooks=[val_init_hook])
all_predicions = []
image_ids = []
num_generated = 0
for batch_result in results:
image_id, pred = batch_result["image_id"], batch_result["predictions"]
result = ''.join([idx_to_word[idx] for idx in pred if idx != 0 and idx != 2])
all_predicions.append(result)
image_ids.append(image_id.decode("utf-8").split(".")[0])
num_generated = num_generated + 1
if num_generated % 1000 == 0:
print("Generated %d" % num_generated)
total_results = [{"image_id": img_id, "caption": pred}
for img_id, pred
in zip(image_ids, all_predicions)]
with open("result.json", "w", encoding="utf-8") as f:
json.dump(total_results, f, ensure_ascii=False)
def __init__(self, id, args):
super().__init__()
self.model_dir = args.model_dir
self.config_fp = os.path.join(self.model_dir, 'bert_config.json')
self.checkpoint_fp = os.path.join(self.model_dir, 'bert_model.ckpt')
self.vocab_fp = os.path.join(args.model_dir, 'vocab.txt')
self.tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_fp)
self.max_seq_len = args.max_seq_len
self.worker_id = id
self.daemon = True
self.model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(self.config_fp),
init_checkpoint=self.checkpoint_fp)
os.environ['CUDA_VISIBLE_DEVICES'] = str(self.worker_id)
self.estimator = Estimator(self.model_fn)
self.dest = None
self._start_t = time.perf_counter()
self.socket = None
self.exit_flag = multiprocessing.Event()
def __init__(self, id, args):
super().__init__()
self.model_dir = args.model_dir
self.config_fp = os.path.join(self.model_dir, 'bert_config.json')
self.checkpoint_fp = os.path.join(self.model_dir, 'bert_model.ckpt')
self.vocab_fp = os.path.join(args.model_dir, 'vocab.txt')
self.tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_fp)
self.max_seq_len = args.max_seq_len
self.worker_id = id
self.daemon = True
self.model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(self.config_fp),
init_checkpoint=self.checkpoint_fp,
pooling_strategy=args.pooling_strategy,
pooling_layer=args.pooling_layer)
os.environ['CUDA_VISIBLE_DEVICES'] = str(self.worker_id)
self.estimator = Estimator(self.model_fn)
self.exit_flag = multiprocessing.Event()
self.logger = set_logger('WORKER-%d' % self.worker_id)
def __init__(self):
# the pooling layer index of bert-original model
self.pooling_layer = [-2]
# the pooling_strategy of bert-original model
self.pooling_strategy = PoolingStrategy.REDUCE_MEAN
# "The maximum total input sequence length after WordPiece tokenization. "
# "Sequences longer than this will be truncated, and sequences shorter "
# "than this will be padded."
self.max_seq_len = 128
self.bert_model_dir = sys_conf["bert_dir"]
self.config_fp = os.path.join(self.bert_model_dir, "bert_config.json")
self.ckpt_fp = os.path.join(self.bert_model_dir, "bert_model.ckpt")
self.vocab_fp = os.path.join(self.bert_model_dir, "vocab.txt")
self.tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_fp)
self.model_fn = model_fn_builder(
bert_config=modeling.BertConfig.from_json_file(self.config_fp),
init_checkpoint=self.ckpt_fp,
pooling_strategy=self.pooling_strategy,
pooling_layer=self.pooling_layer)
self.estimator = Estimator(self.model_fn)
def test_estimator(self):
def model_fn(features, labels, mode, params):
del params
x = core_layers.dense(features, 100)
x = core_layers.dense(x, 100)
x = core_layers.dense(x, 100)
x = core_layers.dense(x, 100)
y = core_layers.dense(x, 1)
loss = losses.mean_squared_error(labels, y)
opt = adam.AdamOptimizer(learning_rate=0.1)
train_op = opt.minimize(
loss, global_step=training_util.get_or_create_global_step())
return EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
def input_fn():
batch_size = 128
return (constant_op.constant(np.random.randn(batch_size, 100),
dtype=dtypes.float32),
constant_op.constant(np.random.randn(batch_size, 1),
dtype=dtypes.float32))
config = RunConfig(model_dir='ram://estimator-0/',
save_checkpoints_steps=1)
estimator = Estimator(config=config, model_fn=model_fn)
estimator.train(input_fn=input_fn, steps=10)
estimator.train(input_fn=input_fn, steps=10)
estimator.train(input_fn=input_fn, steps=10)
estimator.train(input_fn=input_fn, steps=10)
def _Run(self, is_training, use_trt, batch_size, num_epochs, model_dir):
"""Train or evaluate the model.
Args:
is_training: whether to train or evaluate the model. In training mode,
quantization will be simulated where the quantize_and_dequantize_v2 are
placed.
use_trt: if true, use TRT INT8 mode for evaluation, which will perform
real quantization. Otherwise use native TensorFlow which will perform
simulated quantization. Ignored if is_training is True.
batch_size: batch size.
num_epochs: how many epochs to train. Ignored if is_training is False.
model_dir: where to save or load checkpoint.
Returns:
The Estimator evaluation result.
"""
# Get dataset
train_data, test_data = mnist.load_data()
def _PreprocessFn(x, y):
x = math_ops.cast(x, dtypes.float32)
x = array_ops.expand_dims(x, axis=2)
x = 2.0 * (x / 255.0) - 1.0
y = math_ops.cast(y, dtypes.int32)
return x, y
def _EvalInputFn():
mnist_x, mnist_y = test_data
dataset = data.Dataset.from_tensor_slices((mnist_x, mnist_y))
dataset = dataset.apply(
data.experimental.map_and_batch(
map_func=_PreprocessFn,
batch_size=batch_size,
num_parallel_calls=8))
dataset = dataset.repeat(count=1)
iterator = dataset.make_one_shot_iterator()
features, labels = iterator.get_next()
return features, labels
def _TrainInputFn():
mnist_x, mnist_y = train_data
dataset = data.Dataset.from_tensor_slices((mnist_x, mnist_y))
dataset = dataset.shuffle(2 * len(mnist_x))
dataset = dataset.apply(
data.experimental.map_and_batch(
map_func=_PreprocessFn,
batch_size=batch_size,
num_parallel_calls=8))
dataset = dataset.repeat(count=num_epochs)
iterator = dataset.make_one_shot_iterator()
features, labels = iterator.get_next()
return features, labels
def _ModelFn(features, labels, mode):
if is_training:
logits_out = self._BuildGraph(features)
else:
graph_def = self._GetGraphDef(use_trt, batch_size, model_dir)
logits_out = importer.import_graph_def(
graph_def,
input_map={INPUT_NODE_NAME: features},
return_elements=[OUTPUT_NODE_NAME + ':0'],
name='')[0]
loss = losses.sparse_softmax_cross_entropy(
labels=labels, logits=logits_out)
summary.scalar('loss', loss)
classes_out = math_ops.argmax(logits_out, axis=1, name='classes_out')
accuracy = metrics.accuracy(
labels=labels, predictions=classes_out, name='acc_op')
summary.scalar('accuracy', accuracy[1])
if mode == ModeKeys.EVAL:
return EstimatorSpec(
mode, loss=loss, eval_metric_ops={'accuracy': accuracy})
elif mode == ModeKeys.TRAIN:
optimizer = AdamOptimizer(learning_rate=1e-2)
train_op = optimizer.minimize(loss, global_step=get_global_step())
return EstimatorSpec(mode, loss=loss, train_op=train_op)
config_proto = config_pb2.ConfigProto()
config_proto.gpu_options.allow_growth = True
estimator = Estimator(
model_fn=_ModelFn,
model_dir=model_dir if is_training else None,
config=RunConfig(session_config=config_proto))
if is_training:
estimator.train(_TrainInputFn)
results = estimator.evaluate(_EvalInputFn)
logging.info('accuracy: %s', str(results['accuracy']))
return results
