I just made my machine learning code work a few days ago and I would like to know if there's a way to improve my code.
Before I get to the implementation of the tasks at hand, I would like to introduce the design I have in mind.
I first created interfaces containing methods that are the tasks needed to be performed in general. These are the interfaces relevant to the code I'm refactoring.
Network Interface
from abc import ABC, abstractmethod
class Network(ABC):
def __init__(self, network_config):
self.network_config = network_config
super(Network, self).__init__()
@abstractmethod
def train(self, train_config, train_features, validation_features, train_labels, validation_labels):
raise NotImplementedError
@abstractmethod
def test(self, test_features, test_labels):
raise NotImplementedError
@abstractmethod
def predict(self, data_to_be_predicted):
raise NotImplementedError
These are the tasks I think that are performed in all Machine Learning problems. The reason for using a network_config
as the initialization parameter and train_config
as parameters is that the parameters (e.g. number of layers for the network, learning rate, etc.) Machine Learning problems use differ from each other. Some require paremeters for the use of momentum and whatnot. Some don't.
My reason for this design is flexibility such that people can switch from one implementation of this interface to another any time.
ImagePreprocessor Interface
from abc import ABC, abstractmethod
class ImagePreprocessor(ABC):
def __init__(self):
super(ImagePreprocessor, self).__init__()
@staticmethod
@abstractmethod
def resize(images, shape):
raise NotImplementedError
@staticmethod
@abstractmethod
def read(image_paths):
raise NotImplementedError
Since the tasks involves the use of images, I also created an interface that will read and preprocess the images.
DatasetPreparator Interface
class DatasetPreparator(ABC):
@staticmethod
@abstractmethod
def get_dataset_from(dataset_config):
raise NotImplementedError
@staticmethod
@abstractmethod
def split_into_train_validation_and_test_sets(features, validation_size, test_size, labels=None):
raise NotImplementedError
Apparently, it's best to split the dataset into training, validation, and testing. Machine learning problems also require getting the dataset.
Here are the implementations for said interfaces:
TFStackedBidirectionalLstmNetwork
def load_charset(charset_file):
return ''.join([line.rstrip('\n') for line in open(charset_file)])
class TensorflowNetwork(Network):
def __init__(self, network_config):
self.graph = tf.Graph()
with self.graph.as_default():
self.inputs = tf.placeholder(tf.float32, [None, None, network_config.num_features], name="input")
self.labels = tf.sparse_placeholder(tf.int32, name="label")
self.seq_len = tf.placeholder(tf.int32, [None], name="seq_len_input")
logits = self._bidirectional_lstm_layers(
network_config.num_hidden_units,
network_config.num_layers,
network_config.num_classes
)
self.global_step = tf.Variable(0, trainable=False)
self.loss = tf.nn.ctc_loss(labels=self.labels, inputs=logits, sequence_length=self.seq_len)
self.cost = tf.reduce_mean(self.loss)
self.optimizer = tf.train.AdamOptimizer(network_config.learning_rate).minimize(self.cost)
self.decoded, _ = tf.nn.ctc_beam_search_decoder(inputs=logits, sequence_length=self.seq_len, merge_repeated=False)
self.dense_decoded = tf.sparse_to_dense(tf.to_int32(self.decoded[0].indices),
tf.to_int32(self.decoded[0].dense_shape),
tf.to_int32(self.decoded[0].values),
name="output")
self.label_error_rate = tf.reduce_mean(tf.edit_distance(tf.cast(self.decoded[0], tf.int32), self.labels))
tf.summary.scalar('cost', self.cost)
self.merged_summary = tf.summary.merge_all()
super(TensorflowNetwork, self).__init__(network_config)
def train(self, train_config, train_features, validation_features, train_labels=None, validation_labels=None):
logger = logging.getLogger('Training for ocr using BI_LSTM_CTC')
logger.setLevel(logging.INFO)
encoder_decoder = EncoderDecoder()
encoder_decoder.initialize_encode_and_decode_maps_from(load_charset(train_config.charset_file))
encoded_train_labels = []
encoded_val_labels = []
for train_label in train_labels:
encoded_train_labels.append(encoder_decoder.encode(train_label))
for val_label in validation_labels:
encoded_val_labels.append(encoder_decoder.encode(val_label))
train_labels = encoded_train_labels
validation_labels = encoded_val_labels
print('loading train data, please wait---------------------', end=' ')
train_feeder = DataIterator(train_features, train_labels, train_config.batch_size)
print('number of training images: ', train_feeder.get_number_of_examples())
print('loading validation data, please wait---------------------', end=' ')
val_feeder = DataIterator(validation_features, validation_labels, train_config.batch_size)
print('number of validation images: ', val_feeder.get_number_of_examples())
num_train_samples = train_feeder.get_number_of_examples()
num_batches_per_epoch = int(num_train_samples/train_config.batch_size)
config = tf.ConfigProto(log_device_placement=False, allow_soft_placement=False)
with tf.Session(graph=self.graph, config=config) as sess:
sess.run(tf.global_variables_initializer())
log_dir = train_config.log_dir + 'train'
saver = tf.train.Saver(max_to_keep=5)
train_writer = tf.summary.FileWriter(log_dir, sess.graph)
if train_config.is_restore:
ckpt = tf.train.latest_checkpoint(train_config.checkpoint_dir)
if ckpt:
saver.restore(sess, ckpt)
print('restore from the checkpoint{0}'.format(ckpt))
print('=============================begin training=============================')
val_inputs, val_seq_len, val_labels = val_feeder.get_whole_data()
val_feed = {
self.inputs: val_inputs,
self.labels: val_labels,
self.seq_len: val_seq_len
}
for current_epoch in range(train_config.num_epochs):
shuffle_index = np.random.permutation(num_train_samples)
train_cost = 0
start = time.time()
for current_batch_number in range(num_batches_per_epoch):
train_feed = self._get_batch_feed(current_batch_number, shuffle_index, train_feeder)
summary_str, batch_cost, step, _ = sess.run([self.merged_summary, self.cost, self.global_step, self.optimizer], train_feed)
train_cost += batch_cost * train_config.batch_size
train_writer.add_summary(summary_str, step)
if not os.path.isdir(train_config.checkpoint_dir):
os.mkdir(train_config.checkpoint_dir)
logger.info('save the checkpoint of {}'.format(step))
tf.train.write_graph(sess.graph_def, log_dir, 'bi_lstm_ctc_ocr.pbtxt')
saver.save(sess, os.path.join(train_config.checkpoint_dir, 'ocr-model-{}.ckpt'.format(current_epoch)), global_step=step)
if step % train_config.validation_steps == 0:
dense_decoded, last_batch_err = sess.run([self.dense_decoded, self.label_error_rate], val_feed)
avg_train_cost = train_cost/((current_batch_number + 1) * train_config.batch_size)
print("Batch {}/{}, Epoch {}/{}, avg_train_cost: {:.3f}, last_batch_err: {:.3f}, time: {:.3f}"
.format(current_batch_number, num_batches_per_epoch, current_epoch + 1, train_config.num_epochs, avg_train_cost, last_batch_err, time.time() - start))
def _get_batch_feed(self, current_batch_number, shuffle_index, train_feeder):
batch_inputs, batch_seq_len, batch_labels = train_feeder.get_next_batch(current_batch_number,
shuffle_index)
feed = {
self.inputs: batch_inputs,
self.labels: batch_labels,
self.seq_len: batch_seq_len
}
return feed
def test(self, test_features, test_labels=None):
print("I'm testing!")
def predict(self, data_to_be_predicted):
print("I'm predicting!")
def _bidirectional_lstm_layers(self, num_hidden, num_layers, num_classes):
lstm_fw_cells = [tf.contrib.rnn.BasicLSTMCell(num_hidden, forget_bias=1.0) for _ in range(num_layers)]
lstm_bw_cells = [tf.contrib.rnn.BasicLSTMCell(num_hidden, forget_bias=1.0) for _ in range(num_layers)]
outputs, _, _ = tf.contrib.rnn.stack_bidirectional_dynamic_rnn(lstm_fw_cells, lstm_bw_cells,
self.inputs, dtype=tf.float32)
batch_size = tf.shape(self.inputs)[0]
outputs = tf.reshape(outputs, [-1, num_hidden])
W = tf.Variable(tf.truncated_normal([num_hidden, num_classes], stddev=0.1, dtype=tf.float32))
b = tf.Variable(tf.constant(0., shape=[num_classes], dtype=tf.float32))
logits = tf.matmul(outputs, W) + b
logits = tf.reshape(logits, [batch_size, -1, num_classes])
logits = tf.transpose(logits, (1, 0, 2))
return logits
I haven't implemented the testing and predicting codes yet. For now, what I'm trying to refactor is the training part and the initialization of the network. As you can see, I used Tensorflow since that's the library I like using.
CV2ImagePreprocessor
class CV2ImagePreprocessor(ImagePreprocessor):
@staticmethod
def resize(images, shape):
resized_images = []
for image in images:
resized_images.append(cv2.resize(image, shape).swapaxes(0,1))
return resized_images
@staticmethod
def read(image_paths):
images = []
for image_path in image_paths:
images.append(cv2.imread(image_path, 0).astype(np.float32))
return images
As you can see, I used CV2 to read and resize the images.
IAMDatasetPreparator
class IAMDatasetPreparator(DatasetPreparator):
@staticmethod
def get_dataset_from(iam_dataset_config):
image_paths = []
labels = []
with open(iam_dataset_config.labels_file) as f:
labeled_data = f.readlines()
labeled_data = [x.strip() for x in labeled_data]
for example in labeled_data:
example_data = example.split()
if example_data[1] == "ok":
image_path = iam_dataset_config.data_dir + example_data[0] + ".png"
image_paths.append(image_path)
label = example_data[-1]
labels.append(label)
return image_paths, labels
@staticmethod
def split_into_train_validation_and_test_sets(images, validation_size, test_size, labels=None):
train_images, test_images, train_labels, test_labels = train_test_split(
images,
labels,
test_size=validation_size,
random_state=128
)
validation_images, test_images, validation_labels, test_labels = train_test_split(
test_images,
test_labels,
test_size=test_size,
random_state=128
)
return train_images, train_labels, validation_images, validation_labels, test_images, test_labels
The dataset I am working with here is the IAM Offline Handwriting Dataset. I used sklearn to split the dataset.
Here's the script I have for training a model.
network_config = NetworkConfig()
network = TensorflowNetwork(network_config)
dataset_config = DatasetConfig()
image_paths, labels = IAMDatasetPreparator.get_image_paths_and_labels_from(dataset_config)
train_data, train_labels, val_data, val_labels, _, _ = \
IAMDatasetPreparator.split_into_train_validation_and_test_sets(image_paths, 0.5, 0.5, labels)
train_config = TrainConfig()
network.train(train_config, train_data, val_data, train_labels, val_labels)
If you want to try it out for yourself, clone this repository. The working implementation uses Tensroflow so Tensorflow is required to get it running. Simply run dummy_train.py
to train a model for 1 epoch. Do note that I'm not trying to improve the accuracy for this question. I'm just trying to know if there's something I can do to improve my code.