# 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

• 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('-tr', '--train_rate', type=float, default=0.8, help='ratio of training data')
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):
"""
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,
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):

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)

• 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.

• Thank you for your great advice. I'll fix correct the code referring to your suggestions. And I'm sorry for not puting a link of entire code, my codes are in my GitHub repository. github.com/Hayashi-Yudai/ML_models – Y. P May 16 at 0:08