U-Net: A TensorFlow model

Question

What I want to know

I am writing a ML model called U-Net with Python (and TensorFlow). My question is not about the machine learning or Tensorflow, I want to know the best structure of the code.

I write some code as my hobby and I've never written code as work, so I do not know what is the "good" code.

My code

I made two files

• main.py
• model.py

The main.py has functions that load data and pass them to the model in model.py. The model.py has a class that represent the ML model and it contains not only the model structure but also a method to train.

The main process of these code is the training method in the model.py, so in main.py, I call this method from the model.py. I feel that this code structure is not correct, that is, I think the main process (training method) should be contained in the main.py, not in the model.py.

I want some advice how I should (or not) change my code.

main.py

import argparse
import os
import glob
import random
import math

import numpy as np
from PIL import Image

import model

def get_parser():
  """
  Set hyper parameters for training UNet.
  """
  parser = argparse.ArgumentParser()
  parser.add_argument('-e', '--epoch', type=int, default=100)
  parser.add_argument('-lr', '--learning_rate', type=float, default=0.0001)
  parser.add_argument('-tr', '--train_rate', type=float, default=0.8, help='ratio of training data')
  parser.add_argument('-b', '--batch_size', type=int, default=20)
  parser.add_argument('-l2', '--l2', type=float, default=0.05, help='L2 regularization')

  return parser

def load_data(image_dir, seg_dir, n_class, train_val_rate, onehot=True):
  """
  load images and segmented images.
  Parameters
  ==========
  image_dir: string 
    the directory of the raw images.
  seg_dir: string
    the directory of the segmented images.
  n_class: int
    the number of classes
  train_val_rate: float
    the ratio of training data to the validation data
  onehot: bool
  Returns
  =======
  the tuple.((training image, training segmented image), (validation image, validation segmented image))
  training/validation (segmented) images are the list of np.ndarray whose shape is (128, 128, 3) (row image)
  and (128, 128, n_class) (segmented image)
  """
  row_img = []
  segmented_img = []
  images = os.listdir(image_dir)
  random.shuffle(images)
  for idx, img in enumerate(images):
    if img.endswith('.png') or img.endswith('.jpg'):
      split_name = os.path.splitext(img)
      img = Image.open(os.path.join(image_dir, img))
      if seg_dir != '':
        seg = Image.open(os.path.join(seg_dir, split_name[0] + '-seg' + split_name[1]))
        seg = seg.resize((128, 128)) 
        seg = np.asarray(seg, dtype=np.int16)
      else:
        seg = None

      img = img.resize((128, 128))
      img = np.asarray(img, dtype=np.float32)
      img, seg = preprocess(img, seg, n_class, onehot=onehot)
      row_img.append(img)
      segmented_img.append(seg)

  train_data = row_img[:int(len(row_img)*train_val_rate)], segmented_img[:int(len(row_img)*train_val_rate)]
  validation_data = row_img[int(len(row_img) * train_val_rate):], segmented_img[int(len(row_img) * train_val_rate):]

  return train_data, validation_data


def generate_data(input_images, teacher_images, batch_size):
  """
  generate the pair of the raw image and segmented image.
  Parameters
  ========== 
  inputs_images: list of the np.array
  teacher_image: list of the np.array or None
  batch_size: int
  Returns
  =======
  """
  batch_num = math.ceil(len(input_images) / batch_size)
  input_images = np.array_split(input_images, batch_num)
  if np.any(teacher_images == None):
    teacher_images = np.zeros(batch_num)
  else:
    teacher_images = np.array_split(teacher_images, batch_num)

  for i in range(batch_num):
    yield input_images[i], teacher_images[i]


def preprocess(img, seg, n_class, onehot):
  if onehot and seg is not None:
    identity = np.identity(n_class, dtype=np.int16) 
    seg = identity[seg]

  return img / 255.0, seg

if __name__ == '__main__':
  parser = get_parser().parse_args()
  unet = model.UNet(classes=2)
  unet.train(parser)

model.py

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

import main

class UNet:
  def __init__(self, classes):
    self.IMAGE_DIR = './dataset/raw_images'
    self.SEGMENTED_DIR = './dataset/segmented_images'
    self.VALIDATION_DIR = './dataset/validation'
    self.classes = classes
    self.X = tf.placeholder(tf.float32, [None, 128, 128, 3]) 
    self.y = tf.placeholder(tf.int16, [None, 128, 128, self.classes])
    self.is_training = tf.placeholder(tf.bool)

  @staticmethod
  def conv2d(
    inputs, filters, kernel_size=3, activation=tf.nn.relu, l2_reg=None, 
    momentum=0.9, epsilon=0.001, is_training=False,
    ):
    """
    convolutional layer. If the l2_reg is a float number, L2 regularization is imposed.

    Parameters
    ----------
      inputs: tf.Tensor
      filters: Non-zero positive integer
        The number of the filter 
      activation: 
        The activation function. The default is tf.nn.relu
      l2_reg: None or float
        The strengthen of the L2 regularization
      is_training: tf.bool
        The default is False. If True, the batch normalization layer is added.
      momentum: float
        The hyper parameter of the batch normalization layer
      epsilon: float
        The hyper parameter of the batch normalization layer
    Returns
    -------
      layer: tf.Tensor
    """
    regularizer = tf.contrib.layers.l2_regularizer(scale=l2_reg) if l2_reg is not None else None
    layer = tf.layers.conv2d(
      inputs=inputs,
      filters=filters,
      kernel_size=kernel_size,
      padding='SAME',
      activation=activation,
      kernel_regularizer=regularizer
    )

    if is_training is not None:
      layer = tf.layers.batch_normalization(
        inputs=layer,
        axis=-1,
        momentum=momentum,
        epsilon=epsilon,
        center=True,
        scale=True,
        training=is_training
      )

    return layer

  @staticmethod
  def trans_conv(inputs, filters, activation=tf.nn.relu, kernel_size=2, strides=2, l2_reg=None):
    """
    transposed convolution layer.
    Parameters
    ---------- 
      inputs: tf.Tensor
      filters: int 
        the number of the filter
      activation: 
        the activation function. The default function is the ReLu.
      kernel_size: int
        the kernel size. Default = 2
      strides: int
        strides. Default = 2
      l2_reg: None or float 
        the strengthen of the L2 regularization.
    Returns
    -------
      layer: tf.Tensor
    """
    regularizer = tf.contrib.layers.l2_regularizer(scale=l2_reg) if l2_reg is not None else None

    layer = tf.layers.conv2d_transpose(
      inputs=inputs,
      filters=filters,
      kernel_size=kernel_size,
      strides=strides,
      kernel_regularizer=regularizer
    )

    return layer

  @staticmethod
  def pooling(inputs):
    return tf.layers.max_pooling2d(inputs=inputs, pool_size=2, strides=2)


  def UNet(self, is_training, l2_reg=None):
    """
    UNet structure.
    Parameters
    ----------
      l2_reg: None or float
        The strengthen of the L2 regularization.
      is_training: tf.bool
        Whether the session is for training or validation.
    Returns
    -------
      outputs: tf.Tensor
    """
    conv1_1 = self.conv2d(self.X, filters=64, l2_reg=l2_reg, is_training=is_training)
    conv1_2 = self.conv2d(conv1_1, filters=64, l2_reg=l2_reg, is_training=is_training)
    pool1 = self.pooling(conv1_2)

    conv2_1 = self.conv2d(pool1, filters=128, l2_reg=l2_reg, is_training=is_training)
    conv2_2 = self.conv2d(conv2_1, filters=128, l2_reg=l2_reg, is_training=is_training)
    pool2 = self.pooling(conv2_2)

    conv3_1 = self.conv2d(pool2, filters=256, l2_reg=l2_reg, is_training=is_training)
    conv3_2 = self.conv2d(conv3_1, filters=256, l2_reg=l2_reg, is_training=is_training)
    pool3 = self.pooling(conv3_2)

    conv4_1 = self.conv2d(pool3, filters=512, l2_reg=l2_reg, is_training=is_training)
    conv4_2 = self.conv2d(conv4_1, filters=512, l2_reg=l2_reg, is_training=is_training)
    pool4 = self.pooling(conv4_2)

    conv5_1 = self.conv2d(pool4, filters=1024, l2_reg=l2_reg)
    conv5_2 = self.conv2d(conv5_1, filters=1024, l2_reg=l2_reg)
    concat1 = tf.concat([conv4_2, self.trans_conv(conv5_2, filters=512, l2_reg=l2_reg)], axis=3)

    conv6_1 = self.conv2d(concat1, filters=512, l2_reg=l2_reg)
    conv6_2 = self.conv2d(conv6_1, filters=512, l2_reg=l2_reg)
    concat2 = tf.concat([conv3_2, self.trans_conv(conv6_2, filters=256, l2_reg=l2_reg)], axis=3)

    conv7_1 = self.conv2d(concat2, filters=256, l2_reg=l2_reg)
    conv7_2 = self.conv2d(conv7_1, filters=256, l2_reg=l2_reg)
    concat3 = tf.concat([conv2_2, self.trans_conv(conv7_2, filters=128, l2_reg=l2_reg)], axis=3)

    conv8_1 = self.conv2d(concat3, filters=128, l2_reg=l2_reg)
    conv8_2 = self.conv2d(conv8_1, filters=128, l2_reg=l2_reg)
    concat4 = tf.concat([conv1_2, self.trans_conv(conv8_2, filters=64, l2_reg=l2_reg)], axis=3)

    conv9_1 = self.conv2d(concat4, filters=64, l2_reg=l2_reg)
    conv9_2 = self.conv2d(conv9_1, filters=64, l2_reg=l2_reg)
    outputs = self.conv2d(conv9_2, filters=self.classes, kernel_size=1, activation=None)

    return outputs

  def train(self, parser):
    """
    training operation
    argument of this function are given by functions in main.py
    Parameters
    ----------
      parser: 
        the paser that has some options
    """
    epoch = parser.epoch
    l2 = parser.l2
    batch_size = parser.batch_size
    train_val_rate = parser.train_rate

    output = self.UNet(l2_reg=l2, is_training=self.is_training)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.y, logits=output))
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
      train_ops = tf.train.AdamOptimizer(parser.learning_rate).minimize(loss)

    init = tf.global_variables_initializer()
    saver = tf.train.Saver(max_to_keep=100)
    all_train, all_val = main.load_data(self.IMAGE_DIR, self.SEGMENTED_DIR, n_class=2, train_val_rate=train_val_rate)
    with tf.Session() as sess:
      init.run()
      for e in range(epoch):
        data = main.generate_data(*all_train, batch_size)
        val_data = main.generate_data(*all_val, len(all_val[0]))
        for Input, Teacher in data:
          sess.run(train_ops, feed_dict={self.X: Input, self.y: Teacher, self.is_training: True})
          ls = loss.eval(feed_dict={self.X: Input, self.y: Teacher, self.is_training: None})
          for val_Input, val_Teacher in val_data:
            val_loss = loss.eval(feed_dict={self.X: val_Input, self.y: val_Teacher, self.is_training: None})

        print(f'epoch #{e + 1}, loss = {ls}, val loss = {val_loss}')
        if e % 100 == 0:
          saver.save(sess, f"./params/model_{e + 1}epochs.ckpt")

      self.validation(sess, output)

  def validation(self, sess, output):
    val_image = main.load_data(self.VALIDATION_DIR, '', n_class=2, train_val_rate=1)[0]
    data = main.generate_data(*val_image, batch_size=1)
    for Input, _ in data:
      result = sess.run(output, feed_dict={self.X: Input, self.is_training: None}) 
      break

    result = np.argmax(result[0], axis=2)
    ident = np.identity(3, dtype=np.int8)
    result = ident[result]*255

    plt.imshow((Input[0]*255).astype(np.int16))
    plt.imshow(result, alpha=0.2)
    plt.show()

The entire code are in my GitHub repository (https://github.com/Hayashi-Yudai/ML_models)

Answer 1

• This is clean code for a hobbyist. ML code in particular can end up being a procedural mess, just carrying out one imported function after another, but this is quite good.

• It's great that you have docstrings with explanations of the arguments, but the explanations of the functions/methods themselves are a little sparse. The odd comment would help a lot too.

• Try to follow PEP8. It's the standard style for all Python code and is generally followed by all Python programmers (4 spaces for indent, 2 spaces between functions, etc).

• Immediately I see that you're importing model in main and main in model. Imports should follow a hierarchy with no cycles. i.e. if A imports from B imports from C, C should never import from A, but it's ok if A imports directly from C.

• Typically, a file called main, run, execute or whatever will import other worker functions from elsewhere, then handle calling them in the correct way. This often means argument parsing is contained in the main file but not a lot else. Given that you have a lot of logic for generating and sorting data, I would separate this out.

• When comparing a var to None, use the idioms var is None or var is not None.

• Try to use proper paths with os.path rather than just strings like '/dataset/segmented_images'. This means you can run the script from anywhere and still find the files/directories you want.

• You check seg_dir is not an empty string, I would be more explicit and have it either be None or the full path described above. Explicitly assigning something as None removes ambiguity. Let's say for example a path is where you are, does the path being '' mean you've found it or not? If the path is None however, obviously you haven't found it.

• Frequently, your variable names are exactly the same as the keyword arguments. This is fine, except the names are often fairly non-descript: inputs, filters etc. Otherwise, variable names are quite good (concise and clear).

• You enumerate() over images in load_data() but don't use the index, just loop over it with a for loop.

• You don't use glob as an import, nor activation as an argument (in trans_conv()).

• UNet() would be a lot nicer with loops to remove the repeated code and unnecessary variable assignments. I would suggest one loop for pre-training and another for post.

• In load_data(), caching the result of int(len(row_img) * train_val_rate would make for much cleaner code, and would tell you what that index actually represents (with a good variable name).

• You make quite a few assumptions that the code will work as intended, far too hopeful! As an example, say you generate an empty list; on some condition you fill the list; then you assume the list has been filled and loop over it. This is fine if the condition is fulfilled, but if it's not then the program will break. Take a look at load_data() and consider how easy it would be to find if there was a problem with how row_img is (or isn't) being filled.

That's how the code looks; I can't really test how it runs without some sample data and expected in/outputs.