train.py

import tensorflow as tf
import DataGenerator
import time

learning_rate = 1e-5
batch_size = 2
prefetch = 4
no_of_epochs = 10000
smoothing = 0.00001

# placeholder for training mode
is_training = tf.placeholder(tf.bool)

# input data generator
trainTransforms = [
    DataGenerator.RandomFlip(),
    DataGenerator.HistogramMatching(data_dir='train-data', train_size=40, prob=0.5),
    DataGenerator.RandomSmoothing(prob=0.3),
    DataGenerator.RandomNoise(prob=0.5),
    DataGenerator.Normalization()
    ]

valTransforms = [
    DataGenerator.Normalization()
    ]

TrainDataset = DataGenerator.DataGenerator(
    data_dir='train-data',
    transforms=trainTransforms,
    train=True
    )

ValDataset = DataGenerator.DataGenerator(
    data_dir='val-data',
    transforms=valTransforms,
    train=False
    )

trainDataset = TrainDataset.get_dataset()
trainDataset = trainDataset.shuffle(buffer_size=5)
trainDataset = trainDataset.batch(batch_size)
trainDataset = trainDataset.prefetch(prefetch)

valDataset = ValDataset.get_dataset()
valDataset = valDataset.shuffle(buffer_size=5)
valDataset = valDataset.batch(batch_size)
valDataset = valDataset.prefetch(prefetch)

iterator = tf.data.Iterator.from_structure(trainDataset.output_types, trainDataset.output_shapes)

training_init_op = iterator.make_initializer(trainDataset)
validation_init_op = iterator.make_initializer(valDataset)
next_item = iterator.get_next()

# convolution layer
def conv3d(x, no_of_input_channels, no_of_filters, filter_size, strides, padding, name):
    with tf.variable_scope(name) as scope:
        
        initializer = tf.variance_scaling_initializer()
        
        filter_size.extend([no_of_input_channels, no_of_filters])
        weights = tf.Variable(initializer(filter_size), name='weights')
        biases = tf.Variable(initializer([no_of_filters]), name='biases')
        conv = tf.nn.conv3d(x, weights, strides=strides, padding=padding, name=name)
        conv += biases
                
        return conv

# transposed convolution layer
def upsamp(x, no_of_kernels, name):
    with tf.variable_scope(name) as scope:
        upsamp = tf.layers.conv3d_transpose(x, no_of_kernels, [2,2,2], 2, padding='VALID', use_bias=True, reuse=tf.AUTO_REUSE)
        return upsamp

# PReLu layer
def prelu(x, scope=None):
    with tf.variable_scope(name_or_scope=scope, default_name="prelu", reuse=tf.AUTO_REUSE):
        alpha = tf.get_variable("prelu", shape=x.get_shape()[-1], dtype=x.dtype, initializer=tf.constant_initializer(0.1))
        prelu_out = tf.maximum(0.0, x) + alpha * tf.minimum(0.0, x)
        return prelu_out
    
# model graph
def graph_encoder(x):
        
    fine_grained_features = {}
    
    conv1 = conv3d(x,1,16,[3,3,3],[1,1,1,1,1],'SAME','Conv1_1')
    conv1 = conv3d(conv1,16,16,[3,3,3],[1,1,1,1,1],'SAME','Conv1_2')
    conv1 = tf.layers.batch_normalization(conv1, training=is_training)
    conv1 = prelu(conv1,'prelu1')
    
    res1 = tf.add(x,conv1)
    fine_grained_features['res1'] = res1
    
    down1 = conv3d(res1,16,32,[2,2,2],[1,2,2,2,1],'VALID','DownSampling1')
    down1 = prelu(down1,'down_prelu1')
    
    conv2 = conv3d(down1,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv2_1')
    conv2 = conv3d(conv2,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv2_2')
    conv2= tf.layers.batch_normalization(conv2, training=is_training)
    conv2 = prelu(conv2,'prelu2')
    
    conv3 = conv3d(conv2,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv3_1')
    conv3 = conv3d(conv3,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv3_2')
    conv3 = tf.layers.batch_normalization(conv3, training=is_training)
    conv3 = prelu(conv3,'prelu3')
    
    res2 = tf.add(down1,conv3)
    fine_grained_features['res2'] = res2

    down2 = conv3d(res2,32,64,[2,2,2],[1,2,2,2,1],'VALID','DownSampling2')
    down2 = prelu(down2,'down_prelu2')
    
    conv4 = conv3d(down2,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv4_1')
    conv4 = conv3d(conv4,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv4_2')
    conv4 = tf.layers.batch_normalization(conv4, training=is_training)
    conv4 = prelu(conv4,'prelu4')
    
    conv5 = conv3d(conv4,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv5_1')
    conv5 = conv3d(conv5,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv5_2')
    conv5 = tf.layers.batch_normalization(conv5, training=is_training)
    conv5 = prelu(conv5,'prelu5')
    
    conv6 = conv3d(conv5,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv6_1')
    conv6 = conv3d(conv6,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv6_2')
    conv6 = tf.layers.batch_normalization(conv6, training=is_training)
    conv6 = prelu(conv6,'prelu6')
    
    res3 = tf.add(down2,conv6)
    fine_grained_features['res3'] = res3

    down3 = conv3d(res3,64,128,[2,2,2],[1,2,2,2,1],'VALID','DownSampling3')
    down3 = prelu(down3,'down_prelu3')
    
    conv7 = conv3d(down3,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv7_1')
    conv7 = conv3d(conv7,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv7_2')
    conv7 = tf.layers.batch_normalization(conv7, training=is_training)
    conv7 = prelu(conv7,'prelu7')
    
    conv8 = conv3d(conv7,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv8_1')
    conv8 = conv3d(conv8,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv8_2')
    conv8 = tf.layers.batch_normalization(conv8, training=is_training)
    conv8 = prelu(conv8,'prelu8')
    
    conv9 = conv3d(conv8,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv9_1')
    conv9 = conv3d(conv9,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv9_2')
    conv9 = tf.layers.batch_normalization(conv9, training=is_training)
    conv9 = prelu(conv9,'prelu9')
    
    res4 = tf.add(down3,conv9)
    fine_grained_features['res4'] = res4

    down4 = conv3d(res4,128,256,[2,2,2],[1,2,2,2,1],'VALID','DownSampling4')
    down4 = prelu(down4,'down_prelu4')
    
    conv10 = conv3d(down4,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv10_1')
    conv10 = conv3d(conv10,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv10_2')
    conv10 = tf.layers.batch_normalization(conv10, training=is_training)
    conv10 = prelu(conv10,'prelu10')
    
    conv11 = conv3d(conv10,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv11_1')
    conv11 = conv3d(conv11,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv11_2')
    conv11 = tf.layers.batch_normalization(conv11, training=is_training)
    conv11 = prelu(conv11,'prelu11')
    
    conv12 = conv3d(conv11,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv12_1')
    conv12 = conv3d(conv12,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv12_2')
    conv12 = tf.layers.batch_normalization(conv12, training=is_training)
    conv12 = prelu(conv12,'prelu12')
    
    res5 = tf.add(down4,conv12)
    fine_grained_features['res5'] = res5
    
    return fine_grained_features

def graph_decoder(features):
        
    inp = features['res5']
    
    upsamp1 = upsamp(inp,128,'Upsampling1')
    upsamp1 = prelu(upsamp1,'prelu_upsamp1')
    
    concat1 = tf.concat([upsamp1,features['res4']],axis=4)
    
    conv13 = conv3d(concat1,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv13_1')
    conv13 = conv3d(conv13,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv13_2')
    conv13 = tf.layers.batch_normalization(conv13, training=is_training)
    conv13 = prelu(conv13,'prelu13')
    
    conv14 = conv3d(conv13,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv14_1')
    conv14 = conv3d(conv14,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv14_2')
    conv14 = tf.layers.batch_normalization(conv14, training=is_training)
    conv14 = prelu(conv14,'prelu14')
    
    conv15 = conv3d(conv14,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv15_1')
    conv15 = conv3d(conv15,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv15_2')
    conv15 = tf.layers.batch_normalization(conv15, training=is_training)
    conv15 = prelu(conv15,'prelu15')
    
    res6 = tf.add(concat1,conv15)
    
    upsamp2 = upsamp(res6,64,'Upsampling2')
    upsamp2 = prelu(upsamp2,'prelu_upsamp2')
    
    concat2 = tf.concat([upsamp2,features['res3']],axis=4)
    
    conv16 = conv3d(concat2,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv16_1')
    conv16 = conv3d(conv16,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv16_2')
    conv16 = tf.layers.batch_normalization(conv16, training=is_training)
    conv16 = prelu(conv16,'prelu16')
    
    conv17 = conv3d(conv16,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv17_1')
    conv17 = conv3d(conv17,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv17_2')
    conv17 = tf.layers.batch_normalization(conv17, training=is_training)
    conv17 = prelu(conv17,'prelu17')
    
    conv18 = conv3d(conv17,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv18_1')
    conv18 = conv3d(conv18,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv18_2')
    conv18 = tf.layers.batch_normalization(conv18, training=is_training)
    conv18 = prelu(conv18,'prelu18')
    
    res7 = tf.add(concat2,conv18)
    
    upsamp3 = upsamp(res7,32,'Upsampling3')
    upsamp3 = prelu(upsamp3,'prelu_upsamp3')
    
    concat3 = tf.concat([upsamp3,features['res2']],axis=4)
    
    conv19 = conv3d(concat3,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv19_1')
    conv19 = conv3d(conv19,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv19_2')
    conv19 = tf.layers.batch_normalization(conv19, training=is_training)
    conv19 = prelu(conv19,'prelu19')
    
    conv20 = conv3d(conv19,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv20_1')
    conv20 = conv3d(conv20,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv20_2')
    conv20 = tf.layers.batch_normalization(conv20, training=is_training)
    conv20 = prelu(conv20,'prelu20')
    
    res8 = tf.add(concat3,conv20)
    
    upsamp4 = upsamp(res8,16,'Upsampling4')
    upsamp4 = prelu(upsamp4,'prelu_upsamp4')
    
    concat4 = tf.concat([upsamp4,features['res1']],axis=4)
    
    conv21 = conv3d(concat4,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv21_1')
    conv21 = conv3d(conv21,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv21_2')
    conv21 = tf.layers.batch_normalization(conv21, training=is_training)
    conv21 = prelu(conv21,'prelu21')
    
    res9 = tf.add(concat4,conv21)
    
    conv22 = conv3d(res9,32,1,[1,1,1],[1,1,1,1,1],'SAME','Conv22')
    conv22 = tf.nn.sigmoid(conv22,'sigmoid')
    
    return conv22

# loss and optimizer
def model_fn():
    
    features, labels = next_item
        
    features = tf.reshape(features, [-1, 128, 128, 64, 1])
    labels = tf.cast(tf.reshape(labels, [-1, 128, 128, 64, 1]), dtype=tf.float32)
    
    # writing summaries to tensorboard
    tf.summary.image('features', features[:, :, :, 32:33, 0], max_outputs=2,collections=['val'])
    tf.summary.image('labels', labels[:, :, :, 32:33, 0], max_outputs=2,collections=['val'])
    
    labels = tf.reshape(labels, [-1,128*128*64])
    
    encoded = graph_encoder(features)
    decoded = graph_decoder(encoded)

    decoded = tf.reshape(decoded, [-1, 128, 128, 64])
    tf.summary.image('segmentation', decoded[:, :, :, 32:33], max_outputs=2, collections=['val'])
    
    output =  tf.reshape(decoded, [-1,128*128*64])
    
    # dice loss
    cost = tf.reduce_mean(tf.divide(smoothing + tf.multiply(2.0, tf.reduce_sum(output * labels, axis=-1)),
                 tf.add(tf.reduce_sum(output, axis=-1), tf.reduce_sum(labels, axis=-1))))
    
    tf.summary.scalar('training_loss', cost)
    tf.summary.scalar('val_loss', cost,collections=['val'])
    
    # for batchnorm
    with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
        opt = tf.train.AdamOptimizer(learning_rate)

        grads = tf.gradients(1-cost, tf.trainable_variables())
        grads = list(zip(grads, tf.trainable_variables()))

        training_operation = opt.apply_gradients(grads_and_vars=grads)
    
    for grad, var in grads:
        tf.summary.histogram(var.name.replace(':',"_") + '/gradient', grad)
        tf.summary.histogram(var.name.replace(':',"_") , var)
    
    return cost, training_operation

# running the session 
def train():
    with tf.Session() as sess:
        
        cost, opt = model_fn()
        sess.run(tf.global_variables_initializer())

        # merging tensorflow summaries
        merged = tf.summary.merge_all()
        merged_val = tf.summary.merge_all(key = 'val')

        train_writer = tf.summary.FileWriter('event/train',sess.graph)
        val_writer = tf.summary.FileWriter('event/val')
        
        saver = tf.train.Saver()

        for epoch in range(1, no_of_epochs+1):
            start_time = time.time()
            train_loss = []
            examples = 0

            # initializing iterator with training dataset
            sess.run([training_init_op])

            while(True):
                try:
                    # training procedure
                    examples += 1
                    loss, _, summary = sess.run([cost, opt, merged], feed_dict={is_training: True})
                    train_writer.add_summary(summary,epoch)
                    train_loss.append(loss)
                    print('Epoch: {} - ex: {} - loss: {:.6f}'.format(epoch, examples*batch_size, sum(train_loss)/len(train_loss)), end="\r")
                except tf.errors.OutOfRangeError:
                    val_loss = []
                    val_example = 0

                    # initializing iterator with validation dataset
                    sess.run([validation_init_op])

                    while(True):
                        try:
                            val_example += 1
                            loss, summary_l = sess.run([cost, merged_val], feed_dict={is_training: False})
                            val_writer.add_summary(summary_l,epoch)
                            val_loss.append(loss)
                            print('Epoch: {} - ex: {} - val_loss: {:.6f}'.format(epoch, val_example*batch_size, sum(val_loss)/len(val_loss)), end="\r")

                        except tf.errors.OutOfRangeError:
                            break
                    break
            
            print('Epoch: {}/{} - loss: {:.6f} - val_loss: {:.6f} - time: {:.4f}'.format(epoch, no_of_epochs, 
                sum(train_loss)/len(train_loss), sum(val_loss)/len(val_loss), time.time()-start_time))
            
            # saving weights
            if epoch%20==0:
                saver.save(sess, '/temp/weights_epoch_{0}.ckpt'.format(epoch))

if __name__ == '__main__':
    train()