Source code for datasets.brain.test_input_fn

import os
import tensorflow as tf

from sacred import Experiment
from shutil import rmtree
from .input_fn import brain3D_input_fn_ingred, brain_input_fn

layers = tf.contrib.layers

class TestBrainInputFn(tf.test.TestCase):
    """docstring for SerializeTest"""

    @classmethod
    def setUpClass(self):
        self.ex = Experiment("test_brain3D_input_fn",
            ingredients=[brain3D_input_fn_ingred])

        self.config_updates = {
            "model_dir": "datasets/brain/test_model_dir",
            "brain3D_feeder": {
                "record_dir": "datasets/brain/test_records",
                "record_pattern": "brain_{split}_{idx}.tfrecord",
                "samples_per_record": 2,
                "batch_size": 1,
                "n_epochs": 1
            }
        }

        @self.ex.command
        def train(model_dir):
        
            def model_fn(features, labels):
                img = features["image"]

                with tf.contrib.framework.arg_scope([layers.conv3d],
                                                    activation_fn=tf.nn.relu,
                                                    kernel_size=4,
                                                    stride=4,
                                                    padding="SAME"):

                    conv1 = layers.conv3d(img, num_outputs=16) # 32x32x32
                    conv2 = layers.conv3d(conv1, num_outputs=32) # 8x8x8
                    conv3 = layers.conv3d(conv2, num_outputs=64) # 2x2x2

                predictions = layers.fully_connected(layers.flatten(conv3),
                                                     num_outputs=1,
                                                     activation_fn=tf.nn.relu)

                predictions = tf.squeeze(predictions)

                loss = tf.losses.absolute_difference(labels["age"],
                                                     predictions)

                train_op = tf.train.GradientDescentOptimizer(0.001).minimize(
                    loss,
                    global_step=tf.train.get_global_step())

                return tf.estimator.EstimatorSpec(
                    mode=tf.estimator.ModeKeys.TRAIN,
                    loss=loss,
                    train_op=train_op)

            estimator = tf.estimator.Estimator(model_fn, model_dir)

            estimator.train(brain_input_fn())

        self.run_object = self.ex.run("train",
                                      config_updates=self.config_updates)

    @classmethod
    def tearDownClass(self):
        if os.path.exists(self.config_updates["model_dir"]):
            rmtree(self.config_updates["model_dir"])
        if os.path.exists(self.config_updates["brain3D_feeder"]["record_dir"]):
            rmtree(self.config_updates["brain3D_feeder"]["record_dir"])

if __name__ == '__main__':
    tf.test.main()