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
"""Configures and runs distributional-skew UQ experiments on MNIST."""
17

18
from __future__ import absolute_import
19
from __future__ import division
20
from __future__ import print_function
21

22
from absl import logging
23
import tensorflow.compat.v2 as tf
24

25
from uq_benchmark_2019 import array_utils
26
from uq_benchmark_2019 import experiment_utils
27
from uq_benchmark_2019.mnist import data_lib
28
from uq_benchmark_2019.mnist import hparams_lib
29
from uq_benchmark_2019.mnist import models_lib
30
gfile = tf.io.gfile
31

32

33
def get_experiment_config(method, architecture,
34
                          test_level, output_dir=None):
35
  """Returns model and data configs."""
36
  data_opts_list = data_lib.DATA_OPTIONS_LIST
37
  if test_level:
38
    data_opts_list = data_opts_list[:4]
39

40
  model_opts = hparams_lib.get_tuned_model_options(architecture, method,
41
                                                   fake_data=test_level > 1,
42
                                                   fake_training=test_level > 0)
43
  if output_dir:
44
    experiment_utils.record_config(model_opts, output_dir+'/model_options.json')
45
  return model_opts, data_opts_list
46

47

48
def run(method, architecture, output_dir, test_level):
49
  """Trains a model and records its predictions on configured datasets.
50

51
  Args:
52
    method: Name of modeling method (vanilla, dropout, svi, ll_svi).
53
    architecture: Name of DNN architecture (mlp or dropout).
54
    output_dir: Directory to record the trained model and output stats.
55
    test_level: Zero indicates no testing. One indicates testing with real data.
56
      Two is for testing with fake data.
57
  """
58
  fake_data = test_level > 1
59
  gfile.makedirs(output_dir)
60
  model_opts, data_opts_list = get_experiment_config(method, architecture,
61
                                                     test_level=test_level,
62
                                                     output_dir=output_dir)
63

64
  # Separately build dataset[0] with shuffle=True for training.
65
  dataset_train = data_lib.build_dataset(data_opts_list[0], fake_data=fake_data)
66
  dataset_eval = data_lib.build_dataset(data_opts_list[1], fake_data=fake_data)
67
  model = models_lib.build_and_train(model_opts,
68
                                     dataset_train, dataset_eval, output_dir)
69
  logging.info('Saving model to output_dir.')
70
  model.save_weights(output_dir + '/model.ckpt')
71

72
  for idx, data_opts in enumerate(data_opts_list):
73
    dataset = data_lib.build_dataset(data_opts, fake_data=fake_data)
74
    logging.info('Running predictions for dataset #%d', idx)
75
    stats = models_lib.make_predictions(model_opts, model, dataset)
76
    array_utils.write_npz(output_dir, 'stats_%d.npz' % idx, stats)
77
    del stats['logits_samples']
78
    array_utils.write_npz(output_dir, 'stats_small_%d.npz' % idx, stats)
79