def create_model(config):
input_shape = (*config['resolution'][0:3], 2)
x = Input(shape=input_shape)
out = __vnet_level__(x, [16, 32, 32],
config,
remove_last_conv=config['half_res'])
# down-conv
mu = Conv3D(3, kernel_size=3, padding='same')(out)
log_sigma = Conv3D(3, kernel_size=3, padding='same')(out)
sampled_velocity_maps = Lambda(
sampling, name="variationalVelocitySampling")([mu, log_sigma])
z = Concatenate(name='zVariationalLoss')([mu, log_sigma])
grads = sampled_velocity_maps
if config['half_res']:
disp_low = Lambda(
toDisplacements(steps=config['exponentialSteps']))(grads)
# upsample displacement map
disp_upsampled = Lambda(toUpscaleResampled)(disp_low)
# we need to fix displacement vectors which are too small after upsampling
disp = Lambda(lambda dispMap: tf.scalar_mul(2., dispMap),
name="manifold_walk")(disp_upsampled)
else:
disp = Lambda(toDisplacements, name="manifold_walk")(grads)
warped = Lambda(transformVolume, name="img_warp")([x, disp])
warpedAtlas = Lambda(transformAtlas, name="atlas_warp")([x, disp])
loss = [
empty_loss,
cc3D(),
smoothness(config['batchsize']), sampleLoss,
cc3D()
]
lossWeights = [
0.,
# data term / CC
1.0,
# smoothness
0.000002,
# loglikelihood
0.2,
# data term / CC atlas warp
1.0
]
model = Model(
inputs=x,
outputs=[disp, warped, sampled_velocity_maps, z, warpedAtlas])
model.compile(optimizer=Adam(lr=1e-4),
loss=loss,
loss_weights=lossWeights,
metrics=['accuracy'])
return model
def train(model_name, gpu_id):
model_dir = '../models/' + model_name
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
with tf.device(gpu):
model = networks.unet(vol_size, nf_enc, nf_dec)
model.compile(optimizer=Adam(lr=lr), loss=[
losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param])
# model.load_weights('../models/udrnet2/udrnet1_1/120000.h5')
train_example_gen = datagenerators.example_gen(train_vol_names)
zero_flow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3))
for step in xrange(0, n_iterations):
X = train_example_gen.next()[0]
train_loss = model.train_on_batch(
[X, atlas_vol], [atlas_vol, zero_flow])
if not isinstance(train_loss, list):
train_loss = [train_loss]
printLoss(step, 1, train_loss)
if(step % model_save_iter == 0):
model.save(model_dir + '/' + str(step) + '.h5')
def train(model, gpu_id, lr, n_iterations, reg_param, model_save_iter,
load_iter):
model_dir = '/home/ys895/MAS3_Models'
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# UNET filters
nf_enc = [16, 32, 32, 32]
if (model == 'vm1'):
nf_dec = [32, 32, 32, 32, 8, 8, 3]
else:
nf_dec = [32, 32, 32, 32, 32, 16, 16, 3]
with tf.device(gpu):
model = networks.unet(vol_size, nf_enc, nf_dec)
if (load_iter != 0):
model.load_weights('/home/ys895/MAS3_Models/' + str(load_iter) +
'.h5')
model.compile(optimizer=Adam(lr=lr),
loss=[losses.cc3D(),
losses.gradientLoss('l2')],
loss_weights=[1.0, reg_param])
# model.load_weights('../models/udrnet2/udrnet1_1/120000.h5')
# return the data, add one more dimension into the data
train_example_gen = datagenerators.example_gen(train_vol_names)
zero_flow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3))
# In this part, the code inputs the data into the model
# Before this part, the model was set
for step in range(1, n_iterations + 1):
# choose randomly one of the atlas from the atlas_list
rand_num = random.randint(0, list_num - 1)
atlas_vol = atlas_list[rand_num]
#Parameters for training : X(train_vol) ,atlas_vol(atlas) ,zero_flow
X = train_example_gen.__next__()[0]
train_loss = model.train_on_batch([atlas_vol, X], [X, zero_flow])
if not isinstance(train_loss, list):
train_loss = [train_loss]
printLoss(step, 1, train_loss)
if (step % model_save_iter == 0):
model.save(model_dir + '/' + str(load_iter + step) + '.h5')
def train(model,save_name, gpu_id, lr, n_iterations, reg_param, model_save_iter):
model_dir = '../models/' + save_name
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# UNET filters
nf_enc = [16,32,32,32]
if(model == 'vm1'):
nf_dec = [32,32,32,32,8,8,3]
else:
nf_dec = [32,32,32,32,32,16,16,3]
with tf.device(gpu):
model = networks.unet(vol_size, nf_enc, nf_dec)
model.compile(optimizer=Adam(lr=lr), loss=[losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param])
# model.load_weights('../models/udrnet2/udrnet1_1/120000.h5')
zeroflow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3))
for step in range(0, n_iterations):
sub = np.load(train_pairs[step % (noftrain ** 2)][0])
sub = np.reshape(sub, (1,) + sub.shape + (1,))
tmp = np.load(train_pairs[step % (noftrain ** 2)][1])
tmp = np.reshape(tmp, (1,) + tmp.shape + (1,))
train_loss = model.train_on_batch([sub, tmp], [tmp, zeroflow])
printLoss(step, train_loss, keras.get_value(model.optimizer.lr))
if(step % model_save_iter == 0):
model.save(model_dir + '/' + str(step) + '.h5')
if(step % (2*(noftrain ** 2)) == 0 and step > 0):
keras.set_value(model.optimizer.lr, keras.get_value(model.optimizer.lr) / 2)
def train(model,save_name, gpu_id, lr, n_iterations, reg_param, model_save_iter):
model_dir = '../models/' + save_name
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# UNET filters
nf_enc = [16,32,32,32]
if(model == 'vm1'):
nf_dec = [32,32,32,32,8,8,3]
else:
nf_dec = [32,32,32,32,32,16,16,3]
with tf.device(gpu):
model = networks.unet(vol_size, nf_enc, nf_dec)
model.compile(optimizer=Adam(lr=lr), loss=[
losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param])
# model.load_weights('../models/udrnet2/udrnet1_1/120000.h5')
train_example_gen = datagenerators.example_gen(train_vol_names)
zero_flow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3))
for step in range(0, n_iterations):
X = train_example_gen.__next__()[0]
train_loss = model.train_on_batch(
[X, atlas_vol], [atlas_vol, zero_flow])
if not isinstance(train_loss, list):
train_loss = [train_loss]
printLoss(step, 1, train_loss)
if(step % model_save_iter == 0):
model.save(model_dir + '/' + str(step) + '.h5')
def create_model(config):
input_shape = (*config['resolution'][0:3], 2)
x = Input(shape=input_shape)
# dimension of latent space (batch size by latent dim)
n_z = config['encoding_dense']
# encoder
encoder_filters = config['conv_filters']
tlayer = x
for n_filter in encoder_filters:
tlayer = Conv3D(filters=n_filter,
strides=2,
kernel_size=3,
padding='same',
activation='relu')(tlayer)
tlayer = BatchNormalization()(tlayer) if bool(
config.get("batchnorm", False)) else tlayer
tlayer = Flatten()(tlayer)
# dense ReLU layer to mu and sigma
mu = Dense(n_z, activation='linear')(tlayer)
log_sigma = Dense(n_z, activation='linear')(tlayer)
# sampled latent space
z = Lambda(sampling)([mu, log_sigma])
# decoder
downsampled_scales = list(
reversed([2**i for i, _ in enumerate(encoder_filters)]))
lowest_resolution = [
int(l / downsampled_scales[0]) for l in config['resolution'][0:3]
]
lowest_dim = functools.reduce(lambda x, y: x * y, lowest_resolution)
init_decoder_tensor = Dense(lowest_dim, activation='relu')(z)
tlayer = Reshape(target_shape=(*lowest_resolution, 1))(init_decoder_tensor)
for i, (f, n_filter) in enumerate(
list(zip(downsampled_scales, reversed(encoder_filters)))[:-1]):
conditional_downsampled_moving = AveragePooling3D(
pool_size=f, padding='same')(Lambda(
lambda arg: tf.expand_dims(arg[:, :, :, :, 1], axis=4))(x))
conditional_stack = Concatenate()(
[tlayer, conditional_downsampled_moving])
if config['half_res'] and (i + 1) == len(encoder_filters) - 1:
# do not perform final upconv
tlayer = conditional_stack
else:
# upconvolve
#tlayer = Conv3D(n_filter,kernel_size=3,activation='relu',padding='same')(UpSampling3D(size=2)(conditional_stack))
tlayer = Conv3DTranspose(n_filter,
kernel_size=3,
activation='relu',
strides=2,
padding='same')(conditional_stack)
tlayer = Conv3D(encoder_filters[-1],
kernel_size=3,
activation='relu',
padding='same')(tlayer)
down_conv = Conv3D(16, kernel_size=5, activation='relu',
padding='same')(tlayer)
velocity_maps = Conv3D(3,
kernel_size=5,
activation='relu',
padding='same',
name="velocityMap")(down_conv)
if config['half_res']:
disp_low = Lambda(
toDisplacements(steps=config['exponentialSteps']))(velocity_maps)
# upsample displacement map
disp_upsampled = Lambda(toUpscaleResampled)(disp_low)
# we need to fix displacement vectors which are too small after upsampling
disp = Lambda(lambda dispMap: tf.scalar_mul(2., dispMap),
name="manifold_walk")(disp_upsampled)
else:
disp = Lambda(toDisplacements, name="manifold_walk")(velocity_maps)
# TODO: gaussian smoothing
zLoss = Concatenate(name='zVariationalLoss', axis=2)([
Lambda(lambda a: tf.expand_dims(a, axis=2))(mu),
Lambda(lambda a: tf.expand_dims(a, axis=2))(log_sigma)
])
out = Lambda(transformVolume, name="img_warp")([x, disp])
loss = [empty_loss, cc3D(), smoothness(config['batchsize']), sampleLoss]
lossWeights = [0, 1.5, 0.00001, 0.025]
model = Model(inputs=x, outputs=[disp, out, velocity_maps, zLoss])
model.compile(optimizer=Adam(lr=1e-4),
loss=loss,
loss_weights=lossWeights,
metrics=['accuracy'])
return model
def train(model,
model_dir,
gpu_id,
lr,
n_iterations,
reg_param,
model_save_iter,
batch_size=1):
"""
model training function
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param model_dir: the model directory to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param n_iterations: number of training iterations
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param model_save_iter: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
"""
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# GPU handling
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# UNET filters for voxelmorph-1 and voxelmorph-2,
# these are architectures presented in CVPR 2018
nf_enc = [16, 32, 32, 32]
if model == 'vm1':
nf_dec = [32, 32, 32, 32, 8, 8]
else:
nf_dec = [32, 32, 32, 32, 32, 16, 16]
# prepare the model
# in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow]
# in the experiments, we use image_2 as atlas
model = networks.unet(vol_size, nf_enc, nf_dec)
model.compile(optimizer=Adam(lr=lr),
loss=[losses.cc3D(),
losses.gradientLoss('l2')],
loss_weights=[1.0, reg_param])
# if you'd like to initialize the data, you can do it here:
# model.load_weights(os.path.join(model_dir, '120000.h5'))
# prepare data for training
train_example_gen = datagenerators.example_gen(train_vol_names)
zero_flow = np.zeros([batch_size, *vol_size, 3])
# train. Note: we use train_on_batch and design out own print function as this has enabled
# faster development and debugging, but one could also use fit_generator and Keras callbacks.
for step in range(0, n_iterations):
# get data
X = next(train_example_gen)[0]
# train
train_loss = model.train_on_batch([X, atlas_vol],
[atlas_vol, zero_flow])
if not isinstance(train_loss, list):
train_loss = [train_loss]
# print the loss.
print_loss(step, 1, train_loss)
# save model
if step % model_save_iter == 0:
model.save(os.path.join(model_dir, str(step) + '.h5'))
validation_list.append(val_files[ind[0]][:3] + val_files[ind[1]][:3])
validation_list.append(val_files[ind[1]][:3] + val_files[ind[0]][:3])
gen_train = vol_generator2(datapath, train_list, batch_size)
gen_test = vol_generator2(datapath, validation_list, batch_size)
# training
#history = sdn_refine.fit_generator(gen_train, steps_per_epoch = len(train_list)/batch_size, epochs = epochs, use_multiprocessing = True, verbose=1, validation_data = gen_test, validation_steps = len(validation_list)/batch_size)
#loss = history.history['loss']
#val_loss = history.history['val_loss']
for i in range(3):
set_trainable(sdn, True)
print("Weights in sdn is trainable: {}".format(sdn.trainable))
sdn.compile(
loss=[losses.cc3D(), losses.gradientLoss('l2')],
loss_weights=par['loss_weights'],
# metrics = [rec_img_loss, reg_grad],
optimizer=Adam(lr=par['lr'], decay=1e-5))
sdn.fit_generator(gen_train,
steps_per_epoch=len(train_list) / batch_size,
epochs=3,
use_multiprocessing=True,
verbose=1,
validation_data=gen_test,
validation_steps=len(validation_list) / batch_size)
print('\n ith training in sdn done.\n')
set_trainable(sdn, False)
print("Weights in sdn is trainable: {}".format(sdn.trainable))
print("Weights in sdn model of sdn_refine is trainable: {}".format(
sdn_refine.layers[-3].trainable))
def Loss(yTrue, yPred):
return losses.cc3D()(yTrue, yPred) + losses.cc3D()(warped_back, src)
sdn = Model(inputs, [warped, disp_M(inputs)])
#print(sdn.summary())
print(sdn.layers[-1].summary())
#print(sdn.layers)
from losses import cc3D, gradientLoss, NJ_loss
def reg_loss(y_true, y_pred):
return gradientLoss('l2')(
y_true, y_pred) + par['NJ loss'] * NJ_loss(y_true, y_pred)
sdn.compile(loss=[
cc3D(win=[par['cc_size'], par['cc_size'], par['cc_size']]), reg_loss
],
loss_weights=par['loss_weights'],
optimizer=Adam(lr=par['lr'], decay=1e-4))
'''
replace the following if you have a different split
'''
train_files = ['{:03d}'.format(i) for i in range(38, 80)]
val_files = ['{:03d}'.format(i) for i in range(0, 20)]
train_list = []
validation_list = []
from itertools import combinations
for ind in combinations(range(0, len(train_files), 1), 2):
train_list.append(train_files[ind[0]] + train_files[ind[1]])
train_list.append(train_files[ind[1]] + train_files[ind[0]])
