xcaliber-task/train.py at master · brijml/xcaliber-task · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import keras
from keras.models import Sequential, load_model
from keras.layers import Conv2D, Conv2DTranspose, MaxPooling2D, Dense, Flatten
from keras.losses import categorical_crossentropy
from keras.optimizers import RMSprop, SGD
from keras import regularizers
from logger import Logger

import os, argparse
from data_utils import *

INPUT_SHAPE = (27,27,3)

def get_arguments():
	parser = argparse.ArgumentParser(description='Necessary variables.')
	parser.add_argument('--basepath', type=str, default=1, help = 'path to the dataset directory')
	parser.add_argument('--pretrained', type=int, default=1, help = 'Load pretrained model or not(1/0)')
	parser.add_argument('--optimizer', type=str, default='rms', help = 'which optimizer is to be used')
	parser.add_argument('--modelfile', type=str, default="my-model.h5", help = 'path to be given when pretrained is set to 1')
	parser.add_argument('--lr', type=float, default=1e-4, help = 'learning_rate')
	parser.add_argument('--savedir', type=str, help = 'where the model is saved')
	parser.add_argument('--epoch', type=int, default=5, help = 'number of epochs')
	return parser.parse_args()

def seg_model(n_classes):
	model = Sequential()
	model.add(Conv2D(16, (3,3),
		strides=(1, 1), padding='same',
		use_bias=True, activation='relu',
		kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=regularizers.l2(0.01),
		input_shape=INPUT_SHAPE))
	model.add(Conv2D(16, (3,3),
		strides=(1, 1), padding='same',
		use_bias=True, activation='relu',
		kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=regularizers.l2(0.01)))
	model.add(MaxPooling2D(pool_size=(2, 2)))


	model.add(Conv2D(32, (3,3),
		strides=(1, 1), padding='same',
		use_bias=True, activation='relu',
		kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=regularizers.l2(0.01)))
	model.add(Conv2D(32, (3,3),
		strides=(1, 1), padding='same',
		use_bias=True, activation='relu',
		kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=regularizers.l2(0.01)))
	model.add(MaxPooling2D(pool_size=(2, 2)))

	#Fully connected layers
	model.add(Flatten())
	model.add(Dense(64, activation='relu', kernel_regularizer=regularizers.l2(0.01)))
	model.add(Dense(64, activation='relu', kernel_regularizer=regularizers.l2(0.01)))
	model.add(Dense(25, activation='sigmoid', kernel_regularizer=regularizers.l2(0.01)))

	return model


if __name__ == '__main__':
	args = get_arguments()
	base_path = args.basepath
	image_path = os.path.join(base_path, 'image')
	gt_path = os.path.join(base_path, 'groundTruth')
	gt_files = os.listdir(gt_path)
	train_gt, val_gt = validation_split(gt_files, 0.2)
	print("training on {} samples, validating on {} samples".format(len(train_gt), len(val_gt)))

	save_dir = args.savedir
	epoch = args.epoch
	batch_obj = Logger('batch','batch.log','info')
	logger_batch = batch_obj.log()

	if args.pretrained == 0:
		model = seg_model(2)
		if args.optimizer == 'rms':
			optimizer = RMSprop(lr=args.lr)
		elif args.optimizer == 'sgd':
			optimizer = SGD(lr=args.lr, decay=1e-6, momentum=0.9, nesterov=False)
		model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=['accuracy'])
	else:
		model = load_model(args.modelfile)

	epoch_count = 0
	while epoch_count < epoch:
		#perform training
		batch_count = 0
		train_gen_object = load_gen_v2(image_path, gt_path, train_gt)#, args.batchsize, 2)
		for batch_train in train_gen_object:
			x,y = batch_train
			try:
				assert(x.shape[0] == y.shape[0])
				loss, accuracy = model.train_on_batch(x,y)
				logger_batch.info('training loss for epoch_no {} batch_number {} is loss:{}, accuracy:{}'.format(epoch_count, batch_count, loss, accuracy))
				batch_count+=1

			except Exception as e:
				print(e)
				continue

		#perform validation
		batch_count,total_loss = 0,0
		val_gen_object = load_gen_v2(image_path, gt_path, val_gt)#, args.batchsize, 2)
		for batch_val in val_gen_object:
			x,y = batch_val
			try:
				assert(x.shape[0] == y.shape[0])
				loss, accuracy = model.test_on_batch(x,y)
				batch_count+=1
				total_loss+=loss
			except Exception as e:
				print(e)
				continue

		logger_batch.info('validation loss for epoch_no {} is loss:{}, accuracy:{}'.format(epoch_count, (total_loss/batch_count), accuracy))

		filename = 'mymodel'+str(epoch_count)+'.h5'
		model.save(os.path.join(save_dir, filename))
		epoch_count+=1