eval_mnist_model.py

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

r"""Evaluation script."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os.path as osp
import pdb
import sys
import traceback

from absl import flags
from absl import logging
from monty.collections import AttrDict
import sklearn.cluster
import tensorflow as tf

from capsules.configs import data_config
from capsules.configs import model_config
from capsules.eval import cluster_classify
from capsules.eval import collect_results
from capsules.plot import make_tsne_plot
from capsules.train import tools

flags.DEFINE_string('snapshot', '', 'Checkpoint file.')
flags.DEFINE_string('tsne_figure_name', 'tsne.png', 'Filename for the TSNE '
                                                    'figure. It will be saved in the checkpoint folder.')

# These two flags are necessary for model loading. Don't change them!
flags.DEFINE_string('dataset', 'mnist', 'Don\'t change!')
flags.DEFINE_string('model', 'scae', 'Don\'t change!.')


def _collect_results(sess, tensors, dataset, n_batches):
    """Collects some tensors from many batches."""

    to_collect = AttrDict(
        prior_pres=tensors.caps_presence_prob,
        posterior_pres=tensors.posterior_mixing_probs,
        posterior_acc=tensors.posterior_cls_acc,
        prior_acc=tensors.prior_cls_acc,
        label=dataset['label']
    )

    vals = collect_results(sess, to_collect, n_batches)
    vals.posterior_pres = vals.posterior_pres.sum(1)
    return vals


def main(_=None):
    FLAGS = flags.FLAGS  # pylint: disable=invalid-name,redefined-outer-name
    config = FLAGS
    FLAGS.__dict__['config'] = config

    # Build the graph
    with tf.Graph().as_default():

        model_dict = model_config.get(FLAGS)
        data_dict = data_config.get(FLAGS)

        model = model_dict.model
        trainset = data_dict.trainset
        validset = data_dict.validset

        # Optimisation target
        validset = tools.maybe_convert_dataset(validset)
        trainset = tools.maybe_convert_dataset(trainset)

        train_tensors = model(trainset)
        valid_tensors = model(validset)

        sess = tf.Session()
        saver = tf.train.Saver()
        saver.restore(sess, FLAGS.snapshot)

    valid_results = _collect_results(sess, valid_tensors, validset,
                                     10000 // FLAGS.batch_size)

    train_results = _collect_results(sess, train_tensors, trainset,
                                     60000 // FLAGS.batch_size)

    results = AttrDict(train=train_results, valid=valid_results)

    # Linear classification
    print('Linear classification accuracy:')
    for k, v in results.items():
        print('\t{}: prior={:.04f}, posterior={:.04f}'.format(
            k, v.prior_acc.mean(), v.posterior_acc.mean()))

    # Unsupervised classification via clustering
    print('Bipartite matching classification accuracy:')
    for field in 'posterior_pres prior_pres'.split():
        kmeans = sklearn.cluster.KMeans(
            n_clusters=10,
            precompute_distances=True,
            n_jobs=-1,
            max_iter=1000,
        ).fit(results.train[field])

        train_acc = cluster_classify(results.train[field], results.train.label, 10,
                                     kmeans)
        valid_acc = cluster_classify(results.valid[field], results.valid.label, 10,
                                     kmeans)

        print('\t{}: train_acc={:.04f}, valid_acc={:.04f}'.format(field, train_acc,
                                                                  valid_acc))

    checkpoint_folder = osp.dirname(FLAGS.snapshot)
    figure_filename = osp.join(checkpoint_folder, FLAGS.tsne_figure_name)
    print('Savign TSNE plot at "{}"'.format(figure_filename))
    make_tsne_plot(valid_results.posterior_pres, valid_results.label,
                   figure_filename)


if __name__ == '__main__':
    try:
        logging.set_verbosity(logging.INFO)
        tf.app.run()
    except Exception as err:  # pylint: disable=broad-except
        FLAGS = flags.FLAGS

        last_traceback = sys.exc_info()[2]
        traceback.print_tb(last_traceback)
        print(err)
        pdb.post_mortem(last_traceback)