def train_model(self):
if not os.path.exists(self.MODEL_NAME+'_result'): os.mkdir(self.MODEL_NAME+'_result')
if not os.path.exists(self.LOGS_DIR): os.mkdir(self.LOGS_DIR)
if not os.path.exists(self.CKPT_DIR): os.mkdir(self.CKPT_DIR)
if not os.path.exists(self.OUTPUT_DIR): os.mkdir(self.OUTPUT_DIR)
train_set_path = read_data_path(self.TRAIN_IMAGE_PATH, self.TRAIN_LABEL_PATH)
valid_set_path = read_data_path(self.VALID_IMAGE_PATH, self.VALID_LABEL_PATH)
ckpt_save_path = os.path.join(self.CKPT_DIR, self.MODEL_NAME+'_'+str(self.N_BATCH)+'_'+str(self.LEARNING_RATE))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
total_batch = int(len(train_set_path) / self.N_BATCH)
counter = 0
self.saver = tf.train.Saver()
self.writer = tf.summary.FileWriter(self.LOGS_DIR, sess.graph)
for epoch in tqdm(range(self.N_EPOCH)):
total_loss = 0
random.shuffle(train_set_path) # 매 epoch마다 데이터셋 shuffling
random.shuffle(valid_set_path)
for i in range(int(len(train_set_path) / self.N_BATCH)):
# print(i)
batch_xs_path, batch_ys_path = next_batch(train_set_path, self.N_BATCH, i)
batch_xs = read_image(batch_xs_path, [self.RESIZE, self.RESIZE])
batch_ys = read_annotation(batch_ys_path, [self.RESIZE, self.RESIZE])
feed_dict = {self.input_x: batch_xs, self.label_y: batch_ys, self.is_train: True}
_, summary_str ,loss = sess.run([self.optimizer, self.loss_summary, self.loss], feed_dict=feed_dict)
self.writer.add_summary(summary_str, counter)
counter += 1
total_loss += loss
## validation 과정
valid_xs_path, valid_ys_path = next_batch(valid_set_path, 4, 0)
valid_xs = read_image(valid_xs_path, [self.RESIZE, self.RESIZE])
valid_ys = read_annotation(valid_ys_path, [self.RESIZE, self.RESIZE])
valid_pred = sess.run(self.pred, feed_dict={self.input_x: valid_xs, self.label_y: valid_ys, self.is_train:False})
valid_pred = np.squeeze(valid_pred, axis=3)
valid_ys = np.squeeze(valid_ys, axis=3)
## plotting and save figure
img_save_path = self.OUTPUT_DIR + '/' + str(epoch).zfill(3) + '.png'
draw_plot_segmentation(img_save_path, valid_xs, valid_pred, valid_ys)
print('\nEpoch:', '%03d' % (epoch + 1), 'Avg Loss: {:.6}\t'.format(total_loss / total_batch))
self.saver.save(sess, ckpt_save_path+'_'+str(epoch)+'.model', global_step=counter)
self.saver.save(sess, ckpt_save_path+'_'+str(epoch)+'.model', global_step=counter)
print('Finish save model')
def train_model(self):
if not os.path.exists(self.MODEL_NAME+'_result'): os.mkdir(self.MODEL_NAME+'_result')
if not os.path.exists(self.LOGS_DIR): os.mkdir(self.LOGS_DIR)
if not os.path.exists(self.CKPT_DIR): os.mkdir(self.CKPT_DIR)
if not os.path.exists(self.OUTPUT_DIR): os.mkdir(self.OUTPUT_DIR)
train_set_path = read_data_path(self.TRAIN_IMAGE_PATH, self.TRAIN_ANNOT_PATH)
ckpt_save_path = os.path.join(self.CKPT_DIR, self.MODEL_NAME+'_'+str(self.N_BATCH)+'_'+str(self.LEARNING_RATE))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
total_batch = int(len(train_set_path) / self.N_BATCH)
counter = 0
self.saver = tf.train.Saver()
self.writer = tf.summary.FileWriter(self.LOGS_DIR, sess.graph)
for epoch in range(self.N_EPOCH):
total_loss = 0
for i in range(int(len(train_set_path) / self.N_BATCH)):
batch_xs_path, batch_ys_path = next_batch(train_set_path, self.N_BATCH, i)
batch_xs = read_image(batch_xs_path, self.IMAGE_SHAPE[:2])
batch_ys = read_xml(batch_ys_path, self.N_BATCH, self.IMAGE_SHAPE[0], self.GRID_SHAPE[0], self.N_ANCHORS, self.N_CLASSES, self.CLASSES)
feed_dict = {self.input_x: batch_xs, self.label_y:batch_ys, self.is_train: True}
_, summary_str, loss = sess.run([self.optimizer, self.loss_summary, self.loss], feed_dict=feed_dict)
self.writer.add_summary(summary_str, counter)
total_loss += loss
counter += 1
print('Epoch:', '%03d' % (epoch + 1), 'Avg Loss: {:.6}\t'.format(total_loss / total_batch))
self.saver.save(sess, ckpt_save_path+'_'+str(epoch)+'.model', global_step=counter)
self.saver.save(sess, ckpt_save_path+'_'+str(epoch)+'.model', global_step=counter)
print('Finish save model')
