import sys
import os
from glob import glob
import datetime
import tensorflow as tf
import numpy as np
from time import time
from Segmentation.train.utils import setup_gpu, LearningRateUpdate, Metric
from Segmentation.train.reshape import get_mid_slice, get_mid_vol
from Segmentation.train.validation import validate_best_model
from Segmentation.utils.data_loader import read_tfrecord_3d
from Segmentation.utils.visualise_utils import visualise_sample
from Segmentation.utils.losses import dice_loss, tversky_loss, iou_loss
from Segmentation.utils.losses import iou_loss_eval_3d, dice_coef_eval_3d
from Segmentation.utils.losses import dice_loss_weighted_3d, focal_tversky
from Segmentation.model.vnet import VNet
class Train:
def __init__(self,
epochs,
batch_size,
enable_function,
model,
optimizer,
loss_func,
lr_manager,
predict_slice,
metrics,
tfrec_dir='./Data/tfrecords/',
log_dir="logs"):
self.epochs = epochs
self.batch_size = batch_size
self.enable_function = enable_function
self.model = model
self.optimizer = optimizer
self.loss_func = loss_func
self.lr_manager = lr_manager
self.predict_slice = predict_slice
self.metrics = Metric(metrics)
self.tfrec_dir = tfrec_dir
self.log_dir = log_dir
def train_step(self,
x_train,
y_train,
visualise):
with tf.GradientTape() as tape:
predictions = self.model(x_train, training=True)
loss = self.loss_func(y_train, predictions)
grads = tape.gradient(loss, self.model.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
self.metrics.store_metric(y_train, predictions, training=True)
if visualise:
return loss, predictions
return loss, None
def test_step(self,
x_test,
y_test,
visualise):
predictions = self.model(x_test, training=False)
loss = self.loss_func(y_test, predictions)
self.metrics.store_metric(y_test, predictions, training=False)
if visualise:
return loss, predictions
return loss, None
def train_model_loop(self,
train_ds,
valid_ds,
strategy,
multi_class,
visual_save_freq=5,
debug=False,
num_to_visualise=0):
""" Trains 3D model with custom tf loop and MirrorStrategy
"""
def run_train_strategy(x, y, visualise):
total_step_loss, pred = strategy.run(self.train_step, args=(x, y, visualise, ))
return strategy.reduce(
tf.distribute.ReduceOp.SUM, total_step_loss, axis=None), pred
def run_test_strategy(x, y, visualise):
total_step_loss, pred = strategy.run(self.test_step, args=(x, y, visualise, ))
return strategy.reduce(
tf.distribute.ReduceOp.SUM, total_step_loss, axis=None), pred
# TODO(Joe): This needs to be rewritten so that it works with 2D as well
def distributed_train_epoch(train_ds,
epoch,
strategy,
num_to_visualise,
multi_class,
slice_writer,
vol_writer,
visual_save_freq,
predict_slice):
total_loss, num_train_batch = 0.0, 0.0
is_training = True
use_2d = False
for x_train, y_train in train_ds:
visualise = (num_train_batch < num_to_visualise)
loss, pred = run_train_strategy(x_train, y_train, visualise)
loss /= strategy.num_replicas_in_sync
total_loss += loss
if visualise:
num_to_visualise = visualise_sample(x_train, y_train, pred,
num_to_visualise,
slice_writer, vol_writer,
use_2d, epoch, multi_class, predict_slice, is_training)
num_train_batch += 1
return total_loss / num_train_batch
def distributed_test_epoch(valid_ds,
epoch,
strategy,
num_to_visualise,
multi_class,
slice_writer,
vol_writer,
visual_save_freq,
predict_slice):
total_loss, num_test_batch = 0.0, 0.0
is_training = False
use_2d = False
for x_valid, y_valid in valid_ds:
visualise = (num_test_batch < num_to_visualise)
loss, pred = run_test_strategy(x_valid, y_valid, visualise)
loss /= strategy.num_replicas_in_sync
total_loss += loss
if visualise:
num_to_visualise = visualise_sample(x_train, y_train, pred,
num_to_visualise,
slice_writer, vol_writer,
use_2d, epoch, multi_class, predict_slice, is_training)
num_test_batch += 1
return total_loss / num_test_batch
if self.enable_function:
run_train_strategy = tf.function(run_train_strategy)
run_test_strategy = tf.function(run_test_strategy)
# TODO: This whole chunk of code needs to be refactored. Perhaps write it as a function
name = "/" + self.model.name
db = "/debug" if debug else "/test"
mc = "/multi" if multi_class else "/binary"
log_dir_now = self.log_dir + name + db + mc + datetime.datetime.now().strftime("/%Y%m%d/%H%M%S")
train_summary_writer = tf.summary.create_file_writer(log_dir_now + '/train')
test_summary_writer = tf.summary.create_file_writer(log_dir_now + '/val')
test_min_summary_writer = tf.summary.create_file_writer(log_dir_now + '/val_min')
train_img_slice_writer = tf.summary.create_file_writer(log_dir_now + '/train/img/slice')
test_img_slice_writer = tf.summary.create_file_writer(log_dir_now + '/val/img/slice')
train_img_vol_writer = tf.summary.create_file_writer(log_dir_now + '/train/img/vol')
test_img_vol_writer = tf.summary.create_file_writer(log_dir_now + '/val/img/vol')
lr_summary_writer = tf.summary.create_file_writer(log_dir_now + '/lr')
self.metrics.add_metric_summary_writer(log_dir_now)
best_loss = None
for e in range(self.epochs):
self.optimizer.learning_rate = self.lr_manager.update_lr(e)
et0 = time()
train_loss = distributed_train_epoch(train_ds,
e,
strategy,
num_to_visualise,
multi_class,
train_img_slice_writer,
train_img_vol_writer,
visual_save_freq,
self.predict_slice)
with train_summary_writer.as_default():
tf.summary.scalar('epoch_loss', train_loss, step=e)
# distributed_test_epoch(valid_ds,
# e,
# strategy,
# num_to_visualise,
# multi_class,
# test_img_slice_writer,
# test_img_vol_writer,
# visual_save_freq,
# self.predict_slice)
test_loss = distributed_test_epoch(valid_ds,
e,
strategy,
num_to_visualise,
multi_class,
test_img_slice_writer,
test_img_vol_writer,
visual_save_freq,
self.predict_slice)
with test_summary_writer.as_default():
tf.summary.scalar('epoch_loss', test_loss, step=e)
current_lr = self.optimizer.get_config()['learning_rate']
with lr_summary_writer.as_default():
tf.summary.scalar('epoch_lr', current_lr, step=e)
self.metrics.record_metric_to_summary(e)
metric_str = self.metrics.reset_metrics_get_str()
print(f"Epoch {e+1}/{self.epochs} - {time() - et0:.0f}s - loss: {train_loss:.05f} - val_loss: {test_loss:.05f} - lr: {self.optimizer.get_config()['learning_rate']: .06f}" + metric_str)
if best_loss is None:
self.model.save_weights(os.path.join(log_dir_now + f'/best_weights.tf'))
best_loss = test_loss
else:
if test_loss < best_loss:
self.model.save_weights(os.path.join(log_dir_now + f'/best_weights.tf'))
best_loss = test_loss
with test_min_summary_writer.as_default():
tf.summary.scalar('epoch_loss', best_loss, step=e)
return log_dir_now
def load_datasets(batch_size, buffer_size,
tfrec_dir='./Data/tfrecords/',
multi_class=False,
crop_size=144,
depth_crop_size=80,
aug=[],
predict_slice=False,
):
"""
Loads tf records datasets for 3D models.
"""
args = {
'batch_size': batch_size,
'buffer_size': buffer_size,
'multi_class': multi_class,
'use_keras_fit': False,
'crop_size': crop_size,
'depth_crop_size': depth_crop_size,
'aug': aug,
}
train_ds = read_tfrecord_3d(tfrecords_dir=os.path.join(tfrec_dir, 'train_3d/'),
is_training=True, predict_slice=predict_slice, **args)
valid_ds = read_tfrecord_3d(tfrecords_dir=os.path.join(tfrec_dir, 'valid_3d/'),
is_training=False, predict_slice=predict_slice, **args)
return train_ds, valid_ds
def build_model(num_channels, num_classes, name, **kwargs):
"""
Builds standard vnet for 3D
"""
model = VNet(num_channels, num_classes, name=name, **kwargs)
return model
def main(epochs,
name,
log_dir_now=None,
batch_size=2,
val_batch_size=2,
lr=1e-4,
lr_drop=0.9,
lr_drop_freq=5,
lr_warmup=3,
num_to_visualise=2,
num_channels=4,
buffer_size=4,
enable_function=True,
tfrec_dir='./Data/tfrecords/',
multi_class=False,
crop_size=144,
depth_crop_size=80,
aug=[],
debug=False,
predict_slice=False,
tpu=False,
min_lr=1e-7,
custom_loss=None,
**model_kwargs,
):
t0 = time()
if tpu:
tfrec_dir = 'gs://oai-challenge-dataset/tfrecords'
num_classes = 7 if multi_class else 1
metrics = {
'losses': {
'mIoU': [iou_loss, tf.keras.metrics.Mean(), tf.keras.metrics.Mean(), None, None],
'dice': [dice_loss, tf.keras.metrics.Mean(), tf.keras.metrics.Mean(), None, None]
},
}
if multi_class:
metrics['losses']['mIoU-6ch'] = [iou_loss_eval_3d, tf.keras.metrics.Mean(), tf.keras.metrics.Mean(), None, None]
metrics['losses']['dice-6ch'] = [dice_coef_eval_3d, tf.keras.metrics.Mean(), tf.keras.metrics.Mean(), None, None]
train_ds, valid_ds = load_datasets(batch_size, buffer_size, tfrec_dir, multi_class,
crop_size=crop_size, depth_crop_size=depth_crop_size, aug=aug,
predict_slice=predict_slice)
num_gpu = len(tf.config.experimental.list_physical_devices('GPU'))
steps_per_epoch = len(glob(os.path.join(tfrec_dir, 'train_3d/*'))) / (batch_size)
if tpu:
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='pit-tpu')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)
else:
strategy = tf.distribute.MirroredStrategy()
# strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
with strategy.scope():
if custom_loss is None:
loss_func = tversky_loss if multi_class else dice_loss
elif multi_class and custom_loss == "weighted":
loss_func = dice_loss_weighted_3d
elif multi_class and custom_loss == "focal":
loss_func = focal_tversky
else:
raise NotImplementedError(f"Custom loss: {custom_loss} not implemented.")
lr_manager = LearningRateUpdate(lr, lr_drop, lr_drop_freq, warmup=lr_warmup, min_lr=min_lr)
optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
model = build_model(num_channels, num_classes, name, predict_slice=predict_slice, **model_kwargs)
trainer = Train(epochs, batch_size, enable_function,
model, optimizer, loss_func, lr_manager, predict_slice, metrics,
tfrec_dir=tfrec_dir)
train_ds = strategy.experimental_distribute_dataset(train_ds)
valid_ds = strategy.experimental_distribute_dataset(valid_ds)
if log_dir_now is None:
log_dir_now = trainer.train_model_loop(train_ds, valid_ds, strategy, multi_class, debug, num_to_visualise)
train_time = time() - t0
print(f"Train Time: {train_time:.02f}")
t1 = time()
with strategy.scope():
model = build_model(num_channels, num_classes, name, predict_slice=predict_slice, **model_kwargs)
model.load_weights(os.path.join(log_dir_now + f'/best_weights.tf')).expect_partial()
print("Validation for:", log_dir_now)
if not predict_slice:
total_loss, metric_str = validate_best_model(model,
log_dir_now,
val_batch_size,
buffer_size,
tfrec_dir,
multi_class,
crop_size,
depth_crop_size,
predict_slice,
Metric(metrics))
print(f"Train Time: {train_time:.02f}")
print(f"Validation Time: {time() - t1:.02f}")
print(f"Total Time: {time() - t0:.02f}")
with open("results/3d_result.txt", "a") as f:
f.write(f'{log_dir_now}: total_loss {total_loss} {metric_str} \n')
