checkpointファイルがないとエラーが出ます

前提・実現したいこと

http://proceedings.mlr.press/v119/kanade20a.html
↑この論文の実験を再現しようとしています
https://github.com/google-research/google-research/tree/master/cubert
↑こちらにある実験を再現するためのソースコードrun_classifier.pyを実行しています

ckpt.metaからのモデルの読み込みができないので困っているので、どなたか教えていただきたいです。

発生している問題・エラーメッセージ

Traceback (most recent call last):
  File "run_classifier.py", line 1026, in <module>
    tf.app.run()
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_core/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/absl/app.py", line 303, in run
    _run_main(main, args)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "run_classifier.py", line 964, in main
    result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3056, in evaluate
    rendezvous.raise_errors()
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/error_handling.py", line 136, in raise_errors
    six.reraise(typ, value, traceback)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/six.py", line 719, in reraise
    raise value
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3051, in evaluate
    name=name)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 480, in evaluate
    name=name)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 522, in _actual_eval
    return _evaluate()
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 504, in _evaluate
    self._evaluate_build_graph(input_fn, hooks, checkpoint_path))
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1511, in _evaluate_build_graph
    self._call_model_fn_eval(input_fn, self.config))
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1547, in _call_model_fn_eval
    features, labels, ModeKeys.EVAL, config)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 2857, in _call_model_fn
    config)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1149, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3126, in _model_fn
    features, labels, is_export_mode=is_export_mode)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 1663, in call_without_tpu
    return self._call_model_fn(features, labels, is_export_mode=is_export_mode)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 1994, in _call_model_fn
    estimator_spec = self._model_fn(features=features, **kwargs)
  File "run_classifier.py", line 662, in model_fn
    ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
  File "/home/jovyan/work/bert/modeling.py", line 330, in get_assignment_map_from_checkpoint
    init_vars = tf.train.list_variables(init_checkpoint)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_core/python/training/checkpoint_utils.py", line 97, in list_variables
    reader = load_checkpoint(ckpt_dir_or_file)
  File "/opt/conda/envs/cubert_v2/lib/python3.7/site-packages/tensorflow_core/python/training/checkpoint_utils.py", line 65, in load_checkpoint
    "given directory %s" % ckpt_dir_or_file)
ValueError: Couldn't find 'checkpoint' file or checkpoints in given directory exception__epochs_20__pre_trained_epochs_1/model.ckpt-378

該当のソースコード

python
1
2プログラムが長すぎるのでcheckpointを読み込んでいる部分だけこちらに掲載します。
3
4init_checkpoint=$DATA_DIR/exception__epochs_20__pre_trained_epochs_1/model.ckpt-378
5
6def model_fn_builder(bert_config, num_labels, init_checkpoint, learning_rate,
7                     num_train_steps, num_warmup_steps, use_tpu,
8                     use_one_hot_embeddings):
9  """Returns `model_fn` closure for TPUEstimator."""
10
11  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
12    """The `model_fn` for TPUEstimator."""
13
14    tf.logging.info("*** Features ***")
15    for name in sorted(features.keys()):
16      tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))
17
18    input_ids = features["input_ids"]
19    input_mask = features["input_mask"]
20    segment_ids = features["segment_ids"]
21    label_ids = features["label_ids"]
22    is_real_example = None
23    if "is_real_example" in features:
24      is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32)
25    else:
26      is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
27
28    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
29
30    (total_loss, per_example_loss, logits,
31     probabilities) = create_model(bert_config, is_training, input_ids,
32                                   input_mask, segment_ids, label_ids,
33                                   num_labels, use_one_hot_embeddings)
34
35    tvars = tf.trainable_variables()
36    initialized_variable_names = {}
37    scaffold_fn = None
38    if init_checkpoint:
39      (assignment_map, initialized_variable_names
40      ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
41      if use_tpu:
42
43        def tpu_scaffold():
44          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
45          return tf.train.Scaffold()
46
47        scaffold_fn = tpu_scaffold
48      else:
49        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
50
51    tf.logging.info("**** Trainable Variables ****")
52    for var in tvars:
53      init_string = ""
54      if var.name in initialized_variable_names:
55        init_string = ", *INIT_FROM_CKPT*"
56      tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
57                      init_string)
58
59    output_spec = None
60    if mode == tf.estimator.ModeKeys.TRAIN:
61
62      train_op = optimization.create_optimizer(total_loss, learning_rate,
63                                               num_train_steps,
64                                               num_warmup_steps, use_tpu)
65
66      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
67          mode=mode,
68          loss=total_loss,
69          train_op=train_op,
70          scaffold_fn=scaffold_fn)
71      return output_spec
72
73    elif mode == tf.estimator.ModeKeys.EVAL:
74
75      def metric_fn(per_example_loss, label_ids, logits, is_real_example):
76        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
77        accuracy = tf.metrics.accuracy(
78            labels=label_ids, predictions=predictions, weights=is_real_example)
79        mean_per_class_accuracy = tf.metrics.mean_per_class_accuracy(
80            labels=label_ids,
81            predictions=predictions,
82            num_classes=num_labels,
83            weights=is_real_example,
84        )
85        loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)
86
87        return {
88            "eval_accuracy":
89                accuracy,
90            "mean_per_class_accuracy":
91                mean_per_class_accuracy,
92            "eval_loss":
93                loss,
94            "eval_false_negatives":
95                tf.metrics.false_negatives(
96                    labels=label_ids,
97                    predictions=predictions,
98                    weights=is_real_example),
99            "eval_false_positives":
100                tf.metrics.false_positives(
101                    labels=label_ids,
102                    predictions=predictions,
103                    weights=is_real_example),
104            "eval_true_negatives":
105                tf.metrics.true_negatives(
106                    labels=label_ids,
107                    predictions=predictions,
108                    weights=is_real_example),
109            "eval_true_positives":
110                tf.metrics.true_positives(
111                    labels=label_ids,
112                    predictions=predictions,
113                    weights=is_real_example),
114        }
115
116      eval_metrics = (metric_fn,
117                      [per_example_loss, label_ids, logits, is_real_example])
118      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
119          mode=mode,
120          loss=total_loss,
121          eval_metrics=eval_metrics,
122          scaffold_fn=scaffold_fn)
123    else:
124      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
125          mode=mode,
126          predictions={"probabilities": probabilities},
127          scaffold_fn=scaffold_fn)
128    return output_spec
129
130  return model_fn
131

試したこと

init_checkpoint=$DATA_DIR/exception__epochs_20__pre_trained_epochs_1/model.ckpt-378
こちらのファイルの中身は
20200621_Python-exception__epochs_20__pre_trained_epochs_1-model.ckpt-378.data-00000-of-00001、
20200621_Python-exception__epochs_20__pre_trained_epochs_1-model.ckpt-378.index、
20200621_Python-exception__epochs_20__pre_trained_epochs_1-model.ckpt-378.meta　
となっています