def load_checkpoint(sess,
checkpoint_dir,
filename=None,
blacklist=(),
prefix=None):
"""
if `filename` is None, we load last checkpoint, otherwise
we ignore `checkpoint_dir` and load the given checkpoint file.
"""
if filename is None:
# load last checkpoint
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
if ckpt is not None:
filename = ckpt.model_checkpoint_path
else:
checkpoint_dir = os.path.dirname(filename)
vars_ = []
var_names = []
for var in tf.global_variables():
if prefix is None or var.name.startswith(prefix):
name = var.name if prefix is None else var.name[len(prefix) + 1:]
vars_.append(var)
var_names.append(name)
var_file = os.path.join(checkpoint_dir, 'vars.pkl')
if os.path.exists(var_file):
with open(var_file, 'rb') as f:
old_names = pickle.load(f)
else:
old_names = list(var_names)
name_mapping = {}
for name in old_names:
name_ = name
for key, value in variable_mapping:
name_ = re.sub(key, value, name_)
name_mapping[name] = name_
var_names_ = []
for name in var_names:
for key, value in reverse_mapping:
name = re.sub(key, value, name)
var_names_.append(name)
vars_ = dict(zip(var_names_, vars_))
variables = {
old_name[:-2]: vars_[new_name]
for old_name, new_name in name_mapping.items()
if new_name in vars_ and not any(prefix in new_name
for prefix in blacklist)
}
if filename is not None:
utils.log('reading model parameters from {}'.format(filename))
tf.train.Saver(variables).restore(sess, filename)
utils.debug('retrieved parameters ({})'.format(len(variables)))
for var in sorted(variables.values(), key=lambda var: var.name):
utils.debug(' {} {}'.format(var.name, var.get_shape()))
def read_data(self, max_train_size, max_dev_size):
utils.debug('reading training data')
train_set = utils.read_dataset(self.filenames.train,
self.extensions,
self.vocabs,
max_size=max_train_size,
binary_input=self.binary_input,
character_level=self.character_level)
self.batch_iterator = utils.read_ahead_batch_iterator(train_set,
self.batch_size,
read_ahead=10)
utils.debug('reading development data')
dev_sets = [
utils.read_dataset(dev,
self.extensions,
self.vocabs,
max_size=max_dev_size,
binary_input=self.binary_input,
character_level=self.character_level)
for dev in self.filenames.dev
]
# subset of the dev set whose perplexity is periodically evaluated
self.dev_batches = [
utils.get_batches(dev_set, batch_size=self.batch_size, batches=-1)
for dev_set in dev_sets
]
def load_checkpoint(sess, checkpoint_dir, filename, variables):
if filename is not None:
ckpt_file = checkpoint_dir + "/" + filename
utils.log('reading model parameters from {}'.format(ckpt_file))
tf.train.Saver(variables).restore(sess, ckpt_file)
utils.debug('retrieved parameters ({})'.format(len(variables)))
for var in sorted(variables, key=lambda var: var.name):
utils.debug(' {} {}'.format(var.name, var.get_shape()))
def initialize(self,
checkpoints=None,
reset=False,
reset_learning_rate=False,
max_to_keep=1,
keep_every_n_hours=0,
sess=None,
**kwargs):
"""
:param checkpoints: list of checkpoints to load (instead of latest checkpoint)
:param reset: don't load latest checkpoint, reset learning rate and global step
:param reset_learning_rate: reset the learning rate to its initial value
:param max_to_keep: keep this many latest checkpoints at all times
:param keep_every_n_hours: and keep checkpoints every n hours
"""
sess = sess or tf.get_default_session()
if keep_every_n_hours <= 0 or keep_every_n_hours is None:
keep_every_n_hours = float('inf')
self.saver = tf.train.Saver(
max_to_keep=max_to_keep,
keep_checkpoint_every_n_hours=keep_every_n_hours,
sharded=False)
sess.run(tf.global_variables_initializer())
blacklist = ['dropout_keep_prob']
if reset_learning_rate or reset:
blacklist.append('learning_rate')
if reset:
blacklist.append('global_step')
if checkpoints and len(self.models) > 1:
assert len(self.models) == len(checkpoints)
for i, checkpoint in enumerate(checkpoints, 1):
load_checkpoint(sess,
None,
checkpoint,
blacklist=blacklist,
prefix='model_{}'.format(i))
elif checkpoints: # load partial checkpoints
for checkpoint in checkpoints: # checkpoint files to load
load_checkpoint(sess, None, checkpoint, blacklist=blacklist)
elif not reset:
load_checkpoint(sess, self.checkpoint_dir, blacklist=blacklist)
utils.debug('global step: {}'.format(self.global_step.eval()))
utils.debug('baseline step: {}'.format(self.baseline_step.eval()))
def align(self, sess, output=None, align_encoder_id=0, **kwargs):
if self.binary and any(self.binary):
raise NotImplementedError
if len(self.filenames.test) != len(self.extensions):
raise Exception('wrong number of input files')
for line_id, lines in enumerate(utils.read_lines(self.filenames.test)):
token_ids = [
sentence if vocab is None else utils.sentence_to_token_ids(
sentence,
vocab.vocab,
character_level=self.character_level.get(ext)) for ext,
vocab, sentence in zip(self.extensions, self.vocabs, lines)
]
_, weights = self.seq2seq_model.step(sess,
data=[token_ids],
forward_only=True,
align=True,
update_model=False)
trg_vocab = self.trg_vocab[0] # FIXME
trg_token_ids = token_ids[len(self.src_ext)]
trg_tokens = [
trg_vocab.reverse[i]
if i < len(trg_vocab.reverse) else utils._UNK
for i in trg_token_ids
]
weights = weights.squeeze()
max_len = weights.shape[1]
utils.debug(weights)
trg_tokens.append(utils._EOS)
src_tokens = lines[align_encoder_id].split()[:max_len -
1] + [utils._EOS]
output_file = '{}.{}.svg'.format(output, line_id +
1) if output is not None else None
utils.heatmap(src_tokens,
trg_tokens,
weights,
output_file=output_file)
def read_data(self, max_train_size, max_dev_size, read_ahead=10, batch_mode='standard', shuffle=True,
crash_test=False, **kwargs):
utils.debug('reading training data')
self.batch_iterator, self.train_size = utils.get_batch_iterator(
self.filenames.train, self.extensions, self.vocabs, self.batch_size,
max_size=max_train_size, character_level=self.character_level, max_seq_len=self.max_len,
read_ahead=read_ahead, mode=batch_mode, shuffle=shuffle, binary=self.binary, crash_test=crash_test
)
utils.debug('reading development data')
dev_sets = [
utils.read_dataset(dev, self.extensions, self.vocabs, max_size=max_dev_size,
character_level=self.character_level, binary=self.binary)[0]
for dev in self.filenames.dev
]
# subset of the dev set whose loss is periodically evaluated
self.dev_batches = [utils.get_batches(dev_set, batch_size=self.batch_size) for dev_set in dev_sets]
def read_data(self,
max_train_size,
max_dev_size,
read_ahead=10,
batch_mode='standard',
shuffle=True,
**kwargs):
utils.debug('reading training data')
train_set = utils.read_dataset(self.filenames.train,
self.extensions,
self.vocabs,
max_size=max_train_size,
binary_input=self.binary_input,
character_level=self.character_level,
max_seq_len=self.max_input_len)
self.train_size = len(train_set)
self.batch_iterator = utils.read_ahead_batch_iterator(
train_set,
self.batch_size,
read_ahead=read_ahead,
mode=batch_mode,
shuffle=shuffle)
utils.debug('reading development data')
dev_sets = [
utils.read_dataset(dev,
self.extensions,
self.vocabs,
max_size=max_dev_size,
binary_input=self.binary_input,
character_level=self.character_level)
for dev in self.filenames.dev
]
# subset of the dev set whose perplexity is periodically evaluated
self.dev_batches = [
utils.get_batches(dev_set, batch_size=self.batch_size)
for dev_set in dev_sets
]
def load_checkpoint(sess, checkpoint_dir, filename=None, blacklist=()):
""" `checkpoint_dir` should be unique to this model
if `filename` is None, we load last checkpoint, otherwise
we ignore `checkpoint_dir` and load the given checkpoint file.
"""
if filename is None:
# load last checkpoint
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
if ckpt is not None:
filename = ckpt.model_checkpoint_path
else:
checkpoint_dir = os.path.dirname(filename)
var_file = os.path.join(checkpoint_dir, 'vars.pkl')
if os.path.exists(var_file):
with open(var_file, 'rb') as f:
var_names = pickle.load(f)
variables = [
var for var in tf.global_variables() if var.name in var_names
]
else:
variables = tf.global_variables()
# remove variables from blacklist
variables = [
var for var in variables
if not any(prefix in var.name for prefix in blacklist)
]
if filename is not None:
utils.log('reading model parameters from {}'.format(filename))
tf.train.Saver(variables).restore(sess, filename)
utils.debug('retrieved parameters ({})'.format(len(variables)))
for var in variables:
utils.debug(' {} {}'.format(var.name, var.get_shape()))
def __init__(self,
name,
encoders,
decoder,
checkpoint_dir,
learning_rate,
learning_rate_decay_factor,
batch_size,
keep_best=1,
load_embeddings=None,
max_input_len=None,
**kwargs):
super(TranslationModel, self).__init__(name, checkpoint_dir, keep_best,
**kwargs)
self.batch_size = batch_size
self.src_ext = [
encoder.get('ext') or encoder.name for encoder in encoders
]
self.trg_ext = decoder.get('ext') or decoder.name
self.extensions = self.src_ext + [self.trg_ext]
self.max_input_len = max_input_len
encoders_and_decoder = encoders + [decoder]
self.binary_input = [
encoder_or_decoder.binary
for encoder_or_decoder in encoders_and_decoder
]
self.character_level = [
encoder_or_decoder.character_level
for encoder_or_decoder in encoders_and_decoder
]
self.learning_rate = tf.Variable(learning_rate,
trainable=False,
name='learning_rate',
dtype=tf.float32)
self.learning_rate_decay_op = self.learning_rate.assign(
self.learning_rate * learning_rate_decay_factor)
with tf.device('/cpu:0'):
self.global_step = tf.Variable(0,
trainable=False,
name='global_step')
self.filenames = utils.get_filenames(extensions=self.extensions,
**kwargs)
# TODO: check that filenames exist
utils.debug('reading vocabularies')
self._read_vocab()
for encoder_or_decoder, vocab in zip(encoders + [decoder],
self.vocabs):
if encoder_or_decoder.vocab_size <= 0 and vocab is not None:
encoder_or_decoder.vocab_size = len(vocab.reverse)
# this adds an `embedding' attribute to each encoder and decoder
utils.read_embeddings(self.filenames.embeddings, encoders + [decoder],
load_embeddings, self.vocabs)
# main model
utils.debug('creating model {}'.format(name))
self.seq2seq_model = Seq2SeqModel(encoders,
decoder,
self.learning_rate,
self.global_step,
max_input_len=max_input_len,
**kwargs)
self.batch_iterator = None
self.dev_batches = None
self.train_size = None
self.use_sgd = False
def main(args=None):
args = parser.parse_args(args)
# read config file and default config
with open('config/default.yaml') as f:
default_config = utils.AttrDict(yaml.safe_load(f))
with open(args.config) as f:
config = utils.AttrDict(yaml.safe_load(f))
if args.learning_rate is not None:
args.reset_learning_rate = True
# command-line parameters have higher precedence than config file
for k, v in vars(args).items():
if v is not None:
config[k] = v
# set default values for parameters that are not defined
for k, v in default_config.items():
config.setdefault(k, v)
if config.score_function:
config.score_functions = evaluation.name_mapping[config.score_function]
if args.crash_test:
config.max_train_size = 0
if not config.debug:
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # disable TensorFlow's debugging logs
decoding_mode = any(arg is not None for arg in (args.decode, args.eval, args.align))
# enforce parameter constraints
assert config.steps_per_eval % config.steps_per_checkpoint == 0, (
'steps-per-eval should be a multiple of steps-per-checkpoint')
assert decoding_mode or args.train or args.save or args.save_embedding, (
'you need to specify at least one action (decode, eval, align, or train)')
assert not (args.average and args.ensemble)
if args.train and args.purge:
utils.log('deleting previous model')
shutil.rmtree(config.model_dir, ignore_errors=True)
os.makedirs(config.model_dir, exist_ok=True)
# copy config file to model directory
config_path = os.path.join(config.model_dir, 'config.yaml')
if args.train and not os.path.exists(config_path):
with open(args.config) as config_file, open(config_path, 'w') as dest_file:
content = config_file.read()
content = re.sub(r'model_dir:.*?\n', 'model_dir: {}\n'.format(config.model_dir), content,
flags=re.MULTILINE)
dest_file.write(content)
# also copy default config
config_path = os.path.join(config.model_dir, 'default.yaml')
if args.train and not os.path.exists(config_path):
shutil.copy('config/default.yaml', config_path)
# copy source code to model directory
tar_path = os.path.join(config.model_dir, 'code.tar.gz')
if args.train and not os.path.exists(tar_path):
with tarfile.open(tar_path, "w:gz") as tar:
for filename in os.listdir('translate'):
if filename.endswith('.py'):
tar.add(os.path.join('translate', filename), arcname=filename)
logging_level = logging.DEBUG if args.verbose else logging.INFO
# always log to stdout in decoding and eval modes (to avoid overwriting precious train logs)
log_path = os.path.join(config.model_dir, config.log_file)
logger = utils.create_logger(log_path if args.train else None)
logger.setLevel(logging_level)
utils.log('label: {}'.format(config.label))
utils.log('description:\n {}'.format('\n '.join(config.description.strip().split('\n'))))
utils.log(' '.join(sys.argv)) # print command line
try: # print git hash
commit_hash = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode().strip()
utils.log('commit hash {}'.format(commit_hash))
except:
pass
utils.log('tensorflow version: {}'.format(tf.__version__))
# log parameters
utils.debug('program arguments')
for k, v in sorted(config.items(), key=itemgetter(0)):
utils.debug(' {:<20} {}'.format(k, pformat(v)))
if isinstance(config.dev_prefix, str):
config.dev_prefix = [config.dev_prefix]
if config.tasks is not None:
config.tasks = [utils.AttrDict(task) for task in config.tasks]
tasks = config.tasks
else:
tasks = [config]
for task in tasks:
for parameter, value in config.items():
task.setdefault(parameter, value)
task.encoders = [utils.AttrDict(encoder) for encoder in task.encoders]
task.decoders = [utils.AttrDict(decoder) for decoder in task.decoders]
for encoder_or_decoder in task.encoders + task.decoders:
for parameter, value in task.items():
encoder_or_decoder.setdefault(parameter, value)
if args.max_len:
args.max_input_len = args.max_len
if args.max_output_len: # override decoder's max len
task.decoders[0].max_len = args.max_output_len
if args.max_input_len: # override encoder's max len
task.encoders[0].max_len = args.max_input_len
config.checkpoint_dir = os.path.join(config.model_dir, 'checkpoints')
# setting random seeds
if config.seed is None:
config.seed = random.randrange(sys.maxsize)
if config.tf_seed is None:
config.tf_seed = random.randrange(sys.maxsize)
utils.log('python random seed: {}'.format(config.seed))
utils.log('tf random seed: {}'.format(config.tf_seed))
random.seed(config.seed)
tf.set_random_seed(config.tf_seed)
device = None
if config.no_gpu:
device = '/cpu:0'
device_id = None
elif config.gpu_id is not None:
device = '/gpu:{}'.format(config.gpu_id)
device_id = config.gpu_id
else:
device_id = 0
# hide other GPUs so that TensorFlow won't use memory on them
os.environ['CUDA_VISIBLE_DEVICES'] = '' if device_id is None else str(device_id)
utils.log('creating model')
utils.log('using device: {}'.format(device))
with tf.device(device):
if config.weight_scale:
if config.initializer == 'uniform':
initializer = tf.random_uniform_initializer(minval=-config.weight_scale, maxval=config.weight_scale)
else:
initializer = tf.random_normal_initializer(stddev=config.weight_scale)
else:
initializer = None
tf.get_variable_scope().set_initializer(initializer)
# exempt from creating gradient ops
config.decode_only = decoding_mode
if config.tasks is not None:
model = MultiTaskModel(**config)
else:
model = TranslationModel(**config)
# count parameters
# not counting parameters created by training algorithm (e.g. Adam)
variables = [var for var in tf.global_variables() if not var.name.startswith('gradients')]
utils.log('model parameters ({})'.format(len(variables)))
parameter_count = 0
for var in sorted(variables, key=lambda var: var.name):
utils.log(' {} {}'.format(var.name, var.get_shape()))
v = 1
for d in var.get_shape():
v *= d.value
parameter_count += v
utils.log('number of parameters: {:.2f}M'.format(parameter_count / 1e6))
tf_config = tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)
tf_config.gpu_options.allow_growth = config.allow_growth
tf_config.gpu_options.per_process_gpu_memory_fraction = config.mem_fraction
def average_checkpoints(main_sess, sessions):
for var in tf.global_variables():
avg_value = sum(sess.run(var) for sess in sessions) / len(sessions)
main_sess.run(var.assign(avg_value))
with tf.Session(config=tf_config) as sess:
best_checkpoint = os.path.join(config.checkpoint_dir, 'best')
params = {'variable_mapping': config.variable_mapping, 'reverse_mapping': config.reverse_mapping,
'rnn_lm_model_dir': None, 'rnn_mt_model_dir': None,
'rnn_lm_cell_name': None, 'origin_model_ckpt': None}
if config.ensemble and len(config.checkpoints) > 1:
model.initialize(config.checkpoints, **params)
elif config.average and len(config.checkpoints) > 1:
model.initialize(reset=True)
sessions = [tf.Session(config=tf_config) for _ in config.checkpoints]
for sess_, checkpoint in zip(sessions, config.checkpoints):
model.initialize(sess=sess_, checkpoints=[checkpoint], **params)
average_checkpoints(sess, sessions)
elif (not config.checkpoints and decoding_mode and
(os.path.isfile(best_checkpoint + '.index') or os.path.isfile(best_checkpoint + '.index'))):
# in decoding and evaluation mode, unless specified otherwise (by `checkpoints`),
# try to load the best checkpoint
model.initialize([best_checkpoint], **params)
else:
# loads last checkpoint, unless `reset` is true
model.initialize(**config)
if config.output is not None:
dirname = os.path.dirname(config.output)
if dirname:
os.makedirs(dirname, exist_ok=True)
try:
if args.save:
model.save()
elif args.save_embedding:
if config.embedding_output_dir is None:
output_dir = "."
else:
output_dir = config.embedding_output_dir
model.save_embedding(output_dir)
elif args.decode is not None:
if config.align is not None:
config.align = True
model.decode(**config)
elif args.eval is not None:
model.evaluate(on_dev=False, **config)
elif args.align is not None:
model.align(**config)
elif args.train:
model.train(**config)
except KeyboardInterrupt:
sys.exit()
def initialize(self,
checkpoints=None,
reset=False,
reset_learning_rate=False,
max_to_keep=1,
keep_every_n_hours=0,
sess=None,
whitelist=None,
blacklist=None,
**kwargs):
"""
:param checkpoints: list of checkpoints to load (instead of latest checkpoint)
:param reset: don't load latest checkpoint, reset learning rate and global step
:param reset_learning_rate: reset the learning rate to its initial value
:param max_to_keep: keep this many latest checkpoints at all times
:param keep_every_n_hours: and keep checkpoints every n hours
"""
sess = sess or tf.get_default_session()
if keep_every_n_hours <= 0 or keep_every_n_hours is None:
keep_every_n_hours = float('inf')
self.saver = tf.train.Saver(
max_to_keep=max_to_keep,
keep_checkpoint_every_n_hours=keep_every_n_hours,
sharded=False)
sess.run(tf.global_variables_initializer())
# load pre-trained embeddings
for encoder_or_decoder, vocab in zip(self.encoders + self.decoders,
self.vocabs):
if encoder_or_decoder.embedding_file:
utils.log('loading embeddings from: {}'.format(
encoder_or_decoder.embedding_file))
embeddings = {}
with open(encoder_or_decoder.embedding_file) as embedding_file:
for line in embedding_file:
word, vector = line.split(' ', 1)
if word in vocab.vocab:
embeddings[word] = np.array(
list(map(float, vector.split())))
# standardize (mean of 0, std of 0.01)
mean = sum(embeddings.values()) / len(embeddings)
std = np.sqrt(
sum((value - mean)**2
for value in embeddings.values())) / (len(embeddings) -
1)
for key in embeddings:
embeddings[key] = 0.01 * (embeddings[key] - mean) / std
# change TensorFlow variable's value
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
embedding_var = tf.get_variable('embedding_' +
encoder_or_decoder.name)
embedding_value = embedding_var.eval()
for word, i in vocab.vocab.items():
if word in embeddings:
embedding_value[i] = embeddings[word]
sess.run(embedding_var.assign(embedding_value))
if whitelist:
with open(whitelist) as f:
whitelist = list(line.strip() for line in f)
if blacklist:
with open(blacklist) as f:
blacklist = list(line.strip() for line in f)
else:
blacklist = []
blacklist.append('dropout_keep_prob')
if reset_learning_rate or reset:
blacklist.append('learning_rate')
if reset:
blacklist.append('global_step')
params = {
k: kwargs.get(k)
for k in ('variable_mapping', 'reverse_mapping')
}
if checkpoints and len(self.models) > 1:
assert len(self.models) == len(checkpoints)
for i, checkpoint in enumerate(checkpoints, 1):
load_checkpoint(sess,
None,
checkpoint,
blacklist=blacklist,
whitelist=whitelist,
prefix='model_{}'.format(i),
**params)
elif checkpoints: # load partial checkpoints
for checkpoint in checkpoints: # checkpoint files to load
load_checkpoint(sess,
None,
checkpoint,
blacklist=blacklist,
whitelist=whitelist,
**params)
elif not reset:
load_checkpoint(sess,
self.checkpoint_dir,
blacklist=blacklist,
whitelist=whitelist,
**params)
utils.debug('global step: {}'.format(self.global_step.eval()))
utils.debug('baseline step: {}'.format(self.baseline_step.eval()))
def train_step(self,
steps_per_checkpoint,
model_dir,
steps_per_eval=None,
max_steps=0,
max_epochs=0,
eval_burn_in=0,
decay_if_no_progress=None,
decay_after_n_epoch=None,
decay_every_n_epoch=None,
sgd_after_n_epoch=None,
sgd_learning_rate=None,
min_learning_rate=None,
loss_function='xent',
use_baseline=True,
**kwargs):
if min_learning_rate is not None and self.learning_rate.eval(
) < min_learning_rate:
utils.debug('learning rate is too small: stopping')
raise utils.FinishedTrainingException
if 0 < max_steps <= self.global_step.eval(
) or 0 < max_epochs <= self.epoch.eval():
raise utils.FinishedTrainingException
start_time = time.time()
if loss_function == 'reinforce':
step_function = self.seq2seq_model.reinforce_step
else:
step_function = self.seq2seq_model.step
res = step_function(next(self.batch_iterator),
update_model=True,
use_sgd=self.training.use_sgd,
update_baseline=True)
self.training.loss += res.loss
self.training.baseline_loss += getattr(res, 'baseline_loss', 0)
self.training.time += time.time() - start_time
self.training.steps += 1
global_step = self.global_step.eval()
epoch = self.epoch.eval()
if decay_after_n_epoch is not None and self.batch_size * global_step >= decay_after_n_epoch * self.train_size:
if decay_every_n_epoch is not None and (
self.batch_size * (global_step - self.training.last_decay)
>= decay_every_n_epoch * self.train_size):
self.learning_rate_decay_op.eval()
utils.debug(' decaying learning rate to: {:.3g}'.format(
self.learning_rate.eval()))
self.training.last_decay = global_step
if sgd_after_n_epoch is not None and epoch >= sgd_after_n_epoch:
if not self.training.use_sgd:
utils.debug('epoch {}, starting to use SGD'.format(epoch + 1))
self.training.use_sgd = True
if sgd_learning_rate is not None:
self.learning_rate.assign(sgd_learning_rate).eval()
self.training.last_decay = global_step # reset learning rate decay
if steps_per_checkpoint and global_step % steps_per_checkpoint == 0:
loss = self.training.loss / self.training.steps
baseline_loss = self.training.baseline_loss / self.training.steps
step_time = self.training.time / self.training.steps
summary = 'step {} epoch {} learning rate {:.3g} step-time {:.3f} loss {:.3f}'.format(
global_step, epoch + 1, self.learning_rate.eval(), step_time,
loss)
if self.name is not None:
summary = '{} {}'.format(self.name, summary)
if use_baseline and loss_function == 'reinforce':
summary = '{} baseline-loss {:.4f}'.format(
summary, baseline_loss)
utils.log(summary)
if decay_if_no_progress and len(
self.training.losses) >= decay_if_no_progress:
if loss >= max(self.training.losses[:decay_if_no_progress]):
self.learning_rate_decay_op.eval()
self.training.losses.append(loss)
self.training.loss, self.training.time, self.training.steps, self.training.baseline_loss = 0, 0, 0, 0
if steps_per_eval and global_step % steps_per_eval == 0 and 0 <= eval_burn_in <= global_step:
eval_dir = 'eval' if self.name is None else 'eval_{}'.format(
self.name)
eval_output = os.path.join(model_dir, eval_dir)
os.makedirs(eval_output, exist_ok=True)
# if there are several dev files, we define several output files
output = [
os.path.join(eval_output,
'{}.{}.out'.format(prefix, global_step))
for prefix in self.dev_prefix
]
kwargs_ = dict(kwargs)
kwargs_['output'] = output
score, *_ = self.evaluate(on_dev=True, **kwargs_)
self.training.scores.append((global_step, score))
if steps_per_eval and global_step % steps_per_eval == 0:
raise utils.EvalException
elif steps_per_checkpoint and global_step % steps_per_checkpoint == 0:
raise utils.CheckpointException
def __init__(self,
encoders,
decoders,
checkpoint_dir,
learning_rate,
learning_rate_decay_factor,
batch_size,
keep_best=1,
dev_prefix=None,
name=None,
ref_ext=None,
pred_edits=False,
dual_output=False,
binary=None,
truncate_lines=True,
ensemble=False,
checkpoints=None,
beam_size=1,
len_normalization=1,
lexicon=None,
debug=False,
**kwargs):
self.batch_size = batch_size
self.character_level = {}
self.binary = []
self.debug = debug
for encoder_or_decoder in encoders + decoders:
encoder_or_decoder.ext = encoder_or_decoder.ext or encoder_or_decoder.name
self.character_level[
encoder_or_decoder.ext] = encoder_or_decoder.character_level
self.binary.append(encoder_or_decoder.get('binary', False))
self.encoders, self.decoders = encoders, decoders
self.char_output = decoders[0].character_level
self.src_ext = [encoder.ext for encoder in encoders]
self.trg_ext = [decoder.ext for decoder in decoders]
self.extensions = self.src_ext + self.trg_ext
self.ref_ext = ref_ext
if self.ref_ext is not None:
self.binary.append(False)
self.pred_edits = pred_edits
self.dual_output = dual_output
self.dev_prefix = dev_prefix
self.name = name
self.max_input_len = [encoder.max_len for encoder in encoders]
self.max_output_len = [decoder.max_len for decoder in decoders]
self.beam_size = beam_size
if truncate_lines:
self.max_len = None # we let seq2seq.get_batch handle long lines (by truncating them)
else: # the line reader will drop lines that are too long
self.max_len = dict(
zip(self.extensions, self.max_input_len + self.max_output_len))
self.learning_rate = tf.Variable(learning_rate,
trainable=False,
name='learning_rate',
dtype=tf.float32)
self.learning_rate_decay_op = self.learning_rate.assign(
self.learning_rate * learning_rate_decay_factor)
with tf.device('/cpu:0'):
self.global_step = tf.Variable(0,
trainable=False,
name='global_step')
self.baseline_step = tf.Variable(0,
trainable=False,
name='baseline_step')
self.filenames = utils.get_filenames(extensions=self.extensions,
dev_prefix=dev_prefix,
name=name,
ref_ext=ref_ext,
binary=self.binary,
**kwargs)
utils.debug('reading vocabularies')
self.vocabs = None
self.src_vocab, self.trg_vocab = None, None
self.read_vocab()
for encoder_or_decoder, vocab in zip(encoders + decoders, self.vocabs):
if vocab:
if encoder_or_decoder.vocab_size: # reduce vocab size
vocab.reverse[:] = vocab.reverse[:encoder_or_decoder.
vocab_size]
for token, token_id in list(vocab.vocab.items()):
if token_id >= encoder_or_decoder.vocab_size:
del vocab.vocab[token]
else:
encoder_or_decoder.vocab_size = len(vocab.reverse)
utils.debug('creating model')
self.models = []
if ensemble and checkpoints is not None:
for i, _ in enumerate(checkpoints, 1):
with tf.variable_scope('model_{}'.format(i)):
model = Seq2SeqModel(encoders,
decoders,
self.learning_rate,
self.global_step,
name=name,
pred_edits=pred_edits,
dual_output=dual_output,
baseline_step=self.baseline_step,
**kwargs)
self.models.append(model)
self.seq2seq_model = self.models[0]
else:
self.seq2seq_model = Seq2SeqModel(encoders,
decoders,
self.learning_rate,
self.global_step,
name=name,
pred_edits=pred_edits,
dual_output=dual_output,
baseline_step=self.baseline_step,
**kwargs)
self.models.append(self.seq2seq_model)
self.seq2seq_model.create_beam_op(self.models, len_normalization)
self.batch_iterator = None
self.dev_batches = None
self.train_size = None
self.saver = None
self.keep_best = keep_best
self.checkpoint_dir = checkpoint_dir
self.epoch = None
self.training = utils.AttrDict() # used to keep track of training
if lexicon:
with open(lexicon) as lexicon_file:
self.lexicon = dict(line.split() for line in lexicon_file)
else:
self.lexicon = None
def train(self,
sess,
beam_size,
steps_per_checkpoint,
steps_per_eval=None,
eval_output=None,
max_steps=0,
max_epochs=0,
eval_burn_in=0,
decay_if_no_progress=5,
decay_after_n_epoch=None,
decay_every_n_epoch=None,
sgd_after_n_epoch=None,
loss_function='xent',
baseline_steps=0,
reinforce_baseline=True,
reward_function=None,
use_edits=False,
**kwargs):
utils.log('reading training and development data')
self.global_step = 0
for model in self.models:
model.read_data(**kwargs)
# those parameters are used to track the progress of each task
model.loss, model.time, model.steps = 0, 0, 0
model.baseline_loss = 0
model.previous_losses = []
global_step = model.global_step.eval(sess)
model.epoch = model.batch_size * global_step // model.train_size
model.last_decay = global_step
for _ in range(global_step): # read all the data up to this step
next(model.batch_iterator)
self.global_step += global_step
# pre-train baseline
if loss_function == 'reinforce' and baseline_steps > 0 and reinforce_baseline:
utils.log('pre-training baseline')
for model in self.models:
baseline_loss = 0
for step in range(1, baseline_steps + 1):
baseline_loss += model.baseline_step(
sess,
reward_function=reward_function,
use_edits=use_edits)
if step % steps_per_checkpoint == 0:
loss = baseline_loss / steps_per_checkpoint
baseline_loss = 0
utils.log('{} step {} baseline loss {:.4f}'.format(
model.name, step, loss))
utils.log('starting training')
while True:
i = np.random.choice(len(self.models), 1, p=self.ratios)[0]
model = self.models[i]
start_time = time.time()
res = model.train_step(sess,
loss_function=loss_function,
reward_function=reward_function,
use_edits=use_edits)
model.loss += res.loss
if loss_function == 'reinforce':
model.baseline_loss += res.baseline_loss
model.time += time.time() - start_time
model.steps += 1
self.global_step += 1
model_global_step = model.global_step.eval(sess)
epoch = model.batch_size * model_global_step / model.train_size
model.epoch = int(epoch) + 1
if decay_after_n_epoch is not None and epoch >= decay_after_n_epoch:
if decay_every_n_epoch is not None and (
model.batch_size *
(model_global_step - model.last_decay) >=
decay_every_n_epoch * model.train_size):
sess.run(model.learning_rate_decay_op)
utils.debug(' decaying learning rate to: {:.4f}'.format(
model.learning_rate.eval()))
model.last_decay = model_global_step
if sgd_after_n_epoch is not None and epoch >= sgd_after_n_epoch:
if not model.use_sgd:
utils.debug(' epoch {}, starting to use SGD'.format(
model.epoch))
model.use_sgd = True
if steps_per_checkpoint and self.global_step % steps_per_checkpoint == 0:
for model_ in self.models:
if model_.steps == 0:
continue
loss_ = model_.loss / model_.steps
step_time_ = model_.time / model_.steps
if loss_function == 'reinforce':
baseline_loss_ = ' baseline loss {:.4f}'.format(
model_.baseline_loss / model_.steps)
model_.baseline_loss = 0
else:
baseline_loss_ = ''
utils.log(
'{} step {} epoch {} learning rate {:.4f} step-time {:.4f}{} loss {:.4f}'
.format(model_.name,
model_.global_step.eval(sess), model.epoch,
model_.learning_rate.eval(), step_time_,
baseline_loss_, loss_))
if decay_if_no_progress and len(
model_.previous_losses) >= decay_if_no_progress:
if loss_ >= max(
model_.previous_losses[:decay_if_no_progress]):
sess.run(model_.learning_rate_decay_op)
model_.previous_losses.append(loss_)
model_.loss, model_.time, model_.steps = 0, 0, 0
model_.eval_step(sess)
self.save(sess)
if steps_per_eval and self.global_step % steps_per_eval == 0 and 0 <= eval_burn_in <= self.global_step:
score = 0
for ratio, model_ in zip(self.ratios, self.models):
if eval_output is None:
output = None
elif len(model_.filenames.dev) > 1:
# if there are several dev files, we define several output files
# TODO: put dev_prefix into the name of the output file (also in the logging output)
output = [
'{}.{}.{}.{}'.format(eval_output, i + 1,
model_.name,
model_.global_step.eval(sess))
for i in range(len(model_.filenames.dev))
]
else:
output = '{}.{}.{}'.format(
eval_output, model_.name,
model_.global_step.eval(sess))
# kwargs_ = {**kwargs, 'output': output}
kwargs_ = dict(kwargs)
kwargs_['output'] = output
scores_ = model_.evaluate(sess,
beam_size,
on_dev=True,
use_edits=use_edits,
**kwargs_)
score_ = scores_[
0] # in case there are several dev files, only the first one counts
# if there is a main task, pick best checkpoint according to its score
# otherwise use the average score across tasks
if self.main_task is None:
score += ratio * score_
elif model_.name == self.main_task:
score = score_
self.manage_best_checkpoints(self.global_step, score)
if 0 < max_steps <= self.global_step or 0 < max_epochs <= epoch:
utils.log('finished training')
# TODO: save models
return
def main(args=None):
args = parser.parse_args(args)
# read config file and default config
with open('config/default.yaml') as f:
default_config = utils.AttrDict(yaml.safe_load(f))
with open(args.config) as f:
config = utils.AttrDict(yaml.safe_load(f))
if args.learning_rate is not None:
args.reset_learning_rate = True
# command-line parameters have higher precedence than config file
for k, v in vars(args).items():
if v is not None:
config[k] = v
# set default values for parameters that are not defined
for k, v in default_config.items():
config.setdefault(k, v)
# enforce parameter constraints
assert config.steps_per_eval % config.steps_per_checkpoint == 0, (
'steps-per-eval should be a multiple of steps-per-checkpoint')
assert args.decode is not None or args.eval or args.train or args.align, (
'you need to specify at least one action (decode, eval, align, or train)'
)
assert not (args.avg_checkpoints and args.ensemble)
if args.purge:
utils.log('deleting previous model')
shutil.rmtree(config.model_dir, ignore_errors=True)
os.makedirs(config.model_dir, exist_ok=True)
# copy config file to model directory
config_path = os.path.join(config.model_dir, 'config.yaml')
if not os.path.exists(config_path):
shutil.copy(args.config, config_path)
# also copy default config
config_path = os.path.join(config.model_dir, 'default.yaml')
if not os.path.exists(config_path):
shutil.copy('config/default.yaml', config_path)
# copy source code to model directory
tar_path = os.path.join(config.model_dir, 'code.tar.gz')
if not os.path.exists(tar_path):
with tarfile.open(tar_path, "w:gz") as tar:
for filename in os.listdir('translate'):
if filename.endswith('.py'):
tar.add(os.path.join('translate', filename),
arcname=filename)
logging_level = logging.DEBUG if args.verbose else logging.INFO
# always log to stdout in decoding and eval modes (to avoid overwriting precious train logs)
log_path = os.path.join(config.model_dir, config.log_file)
logger = utils.create_logger(log_path if args.train else None)
logger.setLevel(logging_level)
utils.log('label: {}'.format(config.label))
utils.log('description:\n {}'.format('\n '.join(
config.description.strip().split('\n'))))
utils.log(' '.join(sys.argv)) # print command line
try: # print git hash
commit_hash = subprocess.check_output(['git', 'rev-parse',
'HEAD']).decode().strip()
utils.log('commit hash {}'.format(commit_hash))
except:
pass
utils.log('tensorflow version: {}'.format(tf.__version__))
# log parameters
utils.debug('program arguments')
for k, v in sorted(config.items(), key=itemgetter(0)):
utils.debug(' {:<20} {}'.format(k, pformat(v)))
if isinstance(config.dev_prefix, str):
config.dev_prefix = [config.dev_prefix]
if config.tasks is not None:
config.tasks = [utils.AttrDict(task) for task in config.tasks]
tasks = config.tasks
else:
tasks = [config]
for task in tasks:
for parameter, value in config.items():
task.setdefault(parameter, value)
task.encoders = [utils.AttrDict(encoder) for encoder in task.encoders]
task.decoders = [utils.AttrDict(decoder) for decoder in task.decoders]
for encoder_or_decoder in task.encoders + task.decoders:
for parameter, value in task.items():
encoder_or_decoder.setdefault(parameter, value)
device = None
if config.no_gpu:
device = '/cpu:0'
elif config.gpu_id is not None:
device = '/gpu:{}'.format(config.gpu_id)
utils.log('creating model')
utils.log('using device: {}'.format(device))
with tf.device(device):
config.checkpoint_dir = os.path.join(config.model_dir, 'checkpoints')
if config.weight_scale:
if config.initializer == 'uniform':
initializer = tf.random_uniform_initializer(
minval=-config.weight_scale, maxval=config.weight_scale)
else:
initializer = tf.random_normal_initializer(
stddev=config.weight_scale)
else:
initializer = None
tf.get_variable_scope().set_initializer(initializer)
config.decode_only = args.decode is not None or args.eval or args.align # exempt from creating gradient ops
if config.tasks is not None:
model = MultiTaskModel(**config)
else:
model = TranslationModel(**config)
# count parameters
utils.log('model parameters ({})'.format(len(tf.global_variables())))
parameter_count = 0
for var in tf.global_variables():
utils.log(' {} {}'.format(var.name, var.get_shape()))
if not var.name.startswith(
'gradients'
): # not counting parameters created by training algorithm (e.g. Adam)
v = 1
for d in var.get_shape():
v *= d.value
parameter_count += v
utils.log('number of parameters: {:.2f}M'.format(parameter_count / 1e6))
tf_config = tf.ConfigProto(log_device_placement=False,
allow_soft_placement=True)
tf_config.gpu_options.allow_growth = config.allow_growth
tf_config.gpu_options.per_process_gpu_memory_fraction = config.mem_fraction
def average_checkpoints(main_sess, sessions):
for var in tf.global_variables():
avg_value = sum(sess.run(var) for sess in sessions) / len(sessions)
main_sess.run(var.assign(avg_value))
with tf.Session(config=tf_config) as sess:
best_checkpoint = os.path.join(config.checkpoint_dir, 'best')
if ((config.ensemble or config.avg_checkpoints)
and (args.eval or args.decode is not None)
and len(config.checkpoints) > 1):
# create one session for each model in the ensemble
sessions = [tf.Session() for _ in config.checkpoints]
for sess_, checkpoint in zip(sessions, config.checkpoints):
model.initialize(sess_, [checkpoint])
if config.ensemble:
sess = sessions
else:
sess = sessions[0]
average_checkpoints(sess, sessions)
elif (not config.checkpoints
and (args.eval or args.decode is not None or args.align)
and (os.path.isfile(best_checkpoint + '.index')
or os.path.isfile(best_checkpoint + '.index'))):
# in decoding and evaluation mode, unless specified otherwise (by `checkpoints`),
# try to load the best checkpoint)
model.initialize(sess, [best_checkpoint])
else:
# loads last checkpoint, unless `reset` is true
model.initialize(sess, **config)
if args.decode is not None:
model.decode(sess, **config)
elif args.eval:
model.evaluate(sess, on_dev=False, **config)
elif args.align:
model.align(sess, **config)
elif args.train:
try:
model.train(sess=sess, **config)
except (KeyboardInterrupt, utils.FinishedTrainingException):
utils.log('exiting...')
model.save(sess)
sys.exit()
def __init__(self,
encoders,
decoder,
learning_rate,
global_step,
max_gradient_norm,
num_samples=512,
dropout_rate=0.0,
freeze_variables=None,
lm_weight=None,
max_output_len=50,
attention=True,
feed_previous=0.0,
optimizer='sgd',
max_input_len=None,
decode_only=False,
len_normalization=1.0,
**kwargs):
self.lm_weight = lm_weight
self.encoders = encoders
self.decoder = decoder
self.learning_rate = learning_rate
self.global_step = global_step
self.encoder_count = len(encoders)
self.trg_vocab_size = decoder.vocab_size
self.trg_cell_size = decoder.cell_size
self.binary_input = [
encoder.name for encoder in encoders if encoder.binary
]
self.max_output_len = max_output_len
self.max_input_len = max_input_len
self.len_normalization = len_normalization
# if we use sampled softmax, we need an output projection
# sampled softmax only makes sense if we sample less than vocabulary size
if num_samples == 0 or num_samples >= self.trg_vocab_size:
output_projection = None
softmax_loss_function = None
else:
with tf.device('/cpu:0'):
with variable_scope.variable_scope('decoder_{}'.format(
decoder.name)):
w = decoders.get_variable_unsafe(
'proj_w', [self.trg_cell_size, self.trg_vocab_size])
w_t = tf.transpose(w)
b = decoders.get_variable_unsafe('proj_b',
[self.trg_vocab_size])
output_projection = (w, b)
def softmax_loss_function(inputs, labels):
with tf.device('/cpu:0'):
labels = tf.reshape(labels, [-1, 1])
return tf.nn.sampled_softmax_loss(w_t, b, inputs, labels,
num_samples,
self.trg_vocab_size)
if dropout_rate > 0:
self.dropout = tf.Variable(1 - dropout_rate,
trainable=False,
name='dropout_keep_prob')
self.dropout_off = self.dropout.assign(1.0)
self.dropout_on = self.dropout.assign(1 - dropout_rate)
else:
self.dropout = None
self.feed_previous = tf.constant(feed_previous, dtype=tf.float32)
self.encoder_inputs = []
self.encoder_input_length = []
self.extensions = [encoder.name
for encoder in encoders] + [decoder.name]
self.encoder_names = [encoder.name for encoder in encoders]
self.decoder_name = decoder.name
self.extensions = self.encoder_names + [self.decoder_name]
for encoder in self.encoders:
if encoder.binary:
placeholder = tf.placeholder(
tf.float32,
shape=[None, None, encoder.embedding_size],
name='encoder_{}'.format(encoder.name))
else:
placeholder = tf.placeholder(tf.int32,
shape=[None, None],
name='encoder_{}'.format(
encoder.name))
self.encoder_inputs.append(placeholder)
self.encoder_input_length.append(
tf.placeholder(tf.int64,
shape=[None],
name='encoder_{}_length'.format(encoder.name)))
self.decoder_inputs = tf.placeholder(tf.int32,
shape=[None, None],
name='decoder_{}'.format(
self.decoder.name))
self.decoder_input = tf.placeholder(tf.int32,
shape=[None],
name='beam_search_{}'.format(
decoder.name))
self.target_weights = tf.placeholder(tf.float32,
shape=[None, None],
name='weight_{}'.format(
self.decoder.name))
self.targets = tf.placeholder(tf.int32,
shape=[None, None],
name='target_{}'.format(
self.decoder.name))
self.decoder_input_length = tf.placeholder(
tf.int64,
shape=[None],
name='decoder_{}_length'.format(decoder.name))
parameters = dict(encoders=encoders,
decoder=decoder,
dropout=self.dropout,
output_projection=output_projection)
self.attention_states, self.encoder_state = decoders.multi_encoder(
self.encoder_inputs,
encoder_input_length=self.encoder_input_length,
**parameters)
decoder = decoders.attention_decoder if attention else decoders.decoder
self.outputs, self.attention_weights = decoder(
attention_states=self.attention_states,
initial_state=self.encoder_state,
decoder_inputs=self.decoder_inputs,
feed_previous=self.feed_previous,
decoder_input_length=self.decoder_input_length,
**parameters)
self.beam_output, self.beam_tensors = decoders.beam_search_decoder(
decoder_input=self.decoder_input,
attention_states=self.attention_states,
initial_state=self.encoder_state,
**parameters)
self.loss = decoders.sequence_loss(
logits=self.outputs,
targets=self.targets,
weights=self.target_weights,
softmax_loss_function=softmax_loss_function)
if not decode_only:
# gradients and SGD update operation for training the model
if freeze_variables is None:
freeze_variables = []
# compute gradient only for variables that are not frozen
frozen_parameters = [
var.name for var in tf.trainable_variables() if any(
re.match(var_, var.name) for var_ in freeze_variables)
]
if frozen_parameters:
utils.debug('frozen parameters: {}'.format(
', '.join(frozen_parameters)))
params = [
var for var in tf.trainable_variables()
if var.name not in frozen_parameters
]
if optimizer.lower() == 'adadelta':
opt = tf.train.AdadeltaOptimizer(learning_rate=learning_rate)
elif optimizer.lower() == 'adagrad':
opt = tf.train.AdagradOptimizer(learning_rate=learning_rate)
elif optimizer.lower() == 'adam':
opt = tf.train.AdamOptimizer(learning_rate=learning_rate)
else:
opt = tf.train.GradientDescentOptimizer(
learning_rate=learning_rate)
gradients = tf.gradients(self.loss, params)
clipped_gradients, self.gradient_norms = tf.clip_by_global_norm(
gradients, max_gradient_norm)
self.updates = opt.apply_gradients(zip(clipped_gradients, params),
global_step=self.global_step)
def tensor_prod(x, w, b):
shape = tf.shape(x)
x = tf.reshape(x, tf.pack([tf.mul(shape[0], shape[1]), shape[2]]))
x = tf.matmul(x, w) + b
x = tf.reshape(x, tf.pack([shape[0], shape[1], b.get_shape()[0]]))
return x
if output_projection is not None:
w, b = output_projection
self.outputs = tensor_prod(self.outputs, w, b)
self.beam_output = tf.nn.xw_plus_b(self.beam_output, w, b)
self.beam_output = tf.nn.softmax(self.beam_output)
def __init__(self,
encoders,
decoders,
checkpoint_dir,
learning_rate,
learning_rate_decay_factor,
batch_size,
keep_best=1,
dev_prefix=None,
score_function='corpus_scores',
name=None,
ref_ext=None,
pred_edits=False,
dual_output=False,
binary=None,
**kwargs):
self.batch_size = batch_size
self.character_level = {}
self.binary = []
for encoder_or_decoder in encoders + decoders:
encoder_or_decoder.ext = encoder_or_decoder.ext or encoder_or_decoder.name
self.character_level[
encoder_or_decoder.ext] = encoder_or_decoder.character_level
self.binary.append(encoder_or_decoder.get('binary', False))
self.char_output = decoders[0].character_level
self.src_ext = [encoder.ext for encoder in encoders]
self.trg_ext = [decoder.ext for decoder in decoders]
self.extensions = self.src_ext + self.trg_ext
self.ref_ext = ref_ext
if self.ref_ext is not None:
self.binary.append(False)
self.pred_edits = pred_edits
self.dual_output = dual_output
self.dev_prefix = dev_prefix
self.name = name
self.max_input_len = [encoder.max_len for encoder in encoders]
self.max_output_len = [decoder.max_len for decoder in decoders]
self.max_len = dict(
zip(self.extensions, self.max_input_len + self.max_output_len))
self.learning_rate = tf.Variable(learning_rate,
trainable=False,
name='learning_rate',
dtype=tf.float32)
self.learning_rate_decay_op = self.learning_rate.assign(
self.learning_rate * learning_rate_decay_factor)
with tf.device('/cpu:0'):
self.global_step = tf.Variable(0,
trainable=False,
name='global_step')
self.baseline_step = tf.Variable(0,
trainable=False,
name='baseline_step')
self.filenames = utils.get_filenames(extensions=self.extensions,
dev_prefix=dev_prefix,
name=name,
ref_ext=ref_ext,
binary=self.binary,
**kwargs)
utils.debug('reading vocabularies')
self.vocabs = None
self.src_vocab, self.trg_vocab = None, None
self.read_vocab()
for encoder_or_decoder, vocab in zip(encoders + decoders, self.vocabs):
if vocab:
encoder_or_decoder.vocab_size = len(vocab.reverse)
utils.debug('creating model')
self.seq2seq_model = Seq2SeqModel(encoders,
decoders,
self.learning_rate,
self.global_step,
name=name,
pred_edits=pred_edits,
dual_output=dual_output,
baseline_step=self.baseline_step,
**kwargs)
self.batch_iterator = None
self.dev_batches = None
self.train_size = None
self.saver = None
self.keep_best = keep_best
self.checkpoint_dir = checkpoint_dir
self.training = utils.AttrDict() # used to keep track of training
try:
self.reversed_scores = getattr(
evaluation, score_function).reversed # the lower the better
except AttributeError:
self.reversed_scores = False # the higher the better
def load_checkpoint(sess, checkpoint_dir, filename=None, blacklist=()):
"""
if `filename` is None, we load last checkpoint, otherwise
we ignore `checkpoint_dir` and load the given checkpoint file.
"""
if filename is None:
# load last checkpoint
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
if ckpt is not None:
filename = ckpt.model_checkpoint_path
else:
checkpoint_dir = os.path.dirname(filename)
var_file = os.path.join(checkpoint_dir, 'vars.pkl')
def get_variable_by_name(name):
for var in tf.global_variables():
if var.name == name:
return var
return None
if os.path.exists(var_file):
with open(var_file, 'rb') as f:
var_names = pickle.load(f)
variables = {}
for var_name in var_names:
skip = False
for var in tf.global_variables():
name = var.name
for key, value in reverse_mapping:
name = re.sub(key, value, name)
if var_name == name:
variables[var_name] = var
skip = True
break
if skip:
continue
name = var_name
for key, value in variable_mapping:
name = re.sub(key, value, name)
for var in tf.global_variables():
if var.name == name:
variables[var_name] = var
break
else:
variables = {var.name: var for var in tf.global_variables()}
# remove variables from blacklist
# variables = [var for var in variables if not any(prefix in var.name for prefix in blacklist)]
variables = {
name[:-2]: var
for name, var in variables.items()
if not any(prefix in name for prefix in blacklist)
}
if filename is not None:
utils.log('reading model parameters from {}'.format(filename))
tf.train.Saver(variables).restore(sess, filename)
utils.debug('retrieved parameters ({})'.format(len(variables)))
for var in sorted(variables.values(), key=lambda var: var.name):
utils.debug(' {} {}'.format(var.name, var.get_shape()))
