google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Tests for load_model."""
17

18
from absl.testing import absltest
19
from absl.testing import parameterized
20
import numpy as np
21
import tensorflow as tf
22
from loss_functions_transfer import load_model
23

24

25
class LoadModelTest(parameterized.TestCase):
26

27
  @parameterized.parameters(list(load_model.LOSS_HYPERPARAMETERS))
28
  def test_build_and_restore(self, loss_name):
29
    seed = 0
30
    inputs_np = np.zeros((1, 224, 224, 3))
31
    labels_np = np.pad([[1]], ((0, 0), (1000, 0)))
32

33
    with tf.Graph().as_default():
34
      inputs = tf.compat.v1.placeholder(tf.float32, (None, 224, 224, 3))
35
      labels = tf.compat.v1.placeholder(tf.float32, (None, 1001))
36
      loss, endpoints = load_model.build_model_and_compute_loss(
37
          loss_name=loss_name, inputs=inputs, labels=labels,
38
          is_training=False)
39
      with tf.compat.v1.Session() as sess:
40
        load_model.restore_checkpoint(loss_name, seed, sess)
41
        sess.run((loss, endpoints),
42
                 feed_dict={inputs: inputs_np, labels: labels_np})
43

44

45
if __name__ == '__main__':
46
  absltest.main()
47