def train( data_dir, val_data_dir, atlas_file, val_atlas_file, model, model_dir, gpu_id, lr, nb_epochs, reg_param, gama_param, steps_per_epoch, batch_size, load_model_file, data_loss, seg_dir=None, # one file val_seg_dir=None, Sf_file=None, # one file val_Sf_file=None, auxi_label=None, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :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 steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param data_loss: 'mse' or 'ncc :param auxi_label: whether to use auxiliary informmation during the training """ # load atlas from provided files. The atlas we used is 160x192x224. # atlas_file = 'D:/voxel/data/t064.tif' atlas = Image.open(atlas_file) # is a TiffImageFile _size is (628, 690) atlas_vol = np.array(atlas)[ np.newaxis, ..., np.newaxis] # is a ndarray, shape is (1, 690, 628, 1) # new = Image.fromarray(X) new.size is (628, 690) vol_size = atlas_vol.shape[1:-1] # (690, 628) print(vol_size) val_atlas = Image.open( val_atlas_file) # is a TiffImageFile _size is (628, 690) val_atlas_vol = np.array(val_atlas)[ np.newaxis, ..., np.newaxis] # is a ndarray, shape is (1, 690, 628, 1) # new = Image.fromarray(X) new.size is (628, 690) val_vol_size = val_atlas_vol.shape[1:-1] # (690, 628) print(val_vol_size) Sm = Image.open(seg_dir) # is a TiffImageFile _size is (628, 690) Sm_ = np.array(Sm)[np.newaxis, ..., np.newaxis] val_Sm = Image.open(val_seg_dir) # is a TiffImageFile _size is (628, 690) val_Sm_ = np.array(val_Sm)[np.newaxis, ..., np.newaxis] # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. # data_dir = D:/voxel/data/01 train_vol_names = data_dir # glob.glob(os.path.join(data_dir, '*.tif')) # is a list contain file path(name) # random.shuffle(train_vol_names) # shuffle volume list tif assert len(train_vol_names) > 0, "Could not find any training data" val_vol_names = val_data_dir # glob.glob(os.path.join(data_dir, '*.tif')) # is a list contain file path(name) # random.shuffle(train_vol_names) # shuffle volume list tif assert len(val_vol_names) > 0, "Could not find any training data" # 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] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC().loss if Sf_file is not None: Sf = Image.open(Sf_file) Sf_ = np.array(Sf)[np.newaxis, ..., np.newaxis] if val_Sf_file is not None: val_Sf = Image.open(val_Sf_file) val_Sf_ = np.array(val_Sf)[np.newaxis, ..., np.newaxis] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) #gpu = gpu_id # data generator nb_gpus = len(gpu_id.split(',')) # 1 assert np.mod(batch_size, nb_gpus) == 0, \ 'batch_size should be a multiple of the nr. of gpus. ' + \ 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus) train_example_gen = datagenerators.example_gen( train_vol_names, batch_size=batch_size) # it is a list contain a ndarray atlas_vol_bs = np.repeat( atlas_vol, batch_size, axis=0) # is a ndarray, if batch_size is 2, shape is (2, 690, 628, 1) cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) val_example_gen = datagenerators.example_gen( val_vol_names, batch_size=batch_size) # it is a list contain a ndarray val_atlas_vol_bs = np.repeat( val_atlas_vol, batch_size, axis=0) # is a ndarray, if batch_size is 2, shape is (2, 690, 628, 1) val_cvpr2018_gen = datagenerators.cvpr2018_gen(val_example_gen, val_atlas_vol_bs, batch_size=batch_size) # prepare the model with tf.device(gpu): sess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) # 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.cvpr2018_net(vol_size, nf_enc, nf_dec) # load initial weights if load_model_file is not None: print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # if auxi_label is not None: # print('yes') # loss_model= [data_loss, losses.Grad('l2').loss, losses.Lseg()._lseg(Sf_) ] ########################## # loss_weight= [1.0, reg_param, gama_param] # else: loss_model = [ data_loss, losses.Grad(gama_param, Sf_, Sm_, penalty='l2').loss ] # real gama: reg_param*gama_param loss_weight = [1.0, reg_param] # reg_param_tensor = tf.constant(5, dtype=tf.float32) metrics_2 = losses.Grad(gama_param, val_Sf_, val_Sm_, penalty='l2', flag_vali=True).loss # reg_param # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(gpu): # multi-gpu support if nb_gpus > 1: save_callback = nrn_gen.ModelCheckpointParallel(save_file_name) mg_model = multi_gpu_model(model, gpus=nb_gpus) # single-gpu else: save_callback = ModelCheckpoint(save_file_name) mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=loss_model, loss_weights=loss_weight, metrics={'flow': metrics_2}) # fit history = mg_model.fit_generator(cvpr2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, validation_data=val_cvpr2018_gen, validation_steps=1, verbose=2) # plot print('model', mg_model.metrics_names) print('keys()', history.history.keys()) # print(metrics.name) plt.plot(history.history['loss']) # plt.plot(history.history['val_spatial_transformer_1_loss']) plt.title('cvpr_auxi_loss') plt.ylabel('loss') plt.xlabel('Epoch') plt.legend(['Train', 'Validation']) plt.show()
def train(data_dir, atlas_file, model, model_name, gpu_id, lr, nb_epochs, reg_param, steps_per_epoch, batch_size, load_model_file, data_loss, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :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 steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param data_loss: data_loss: 'mse' or 'ncc """ # load atlas from provided files. The atlas we used is 160x192x224. # atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] atlas_vol = nib.load(atlas_file).get_data()[np.newaxis, ..., np.newaxis] vol_size = atlas_vol.shape[1:-1] # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. train_vol_names = glob.glob(os.path.join(data_dir, '*.npz')) random.shuffle(train_vol_names) # shuffle volume list assert len(train_vol_names) > 0, "Could not find any training data" # 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] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC().loss model_dir = "../models/" + model_name # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling gpu = '/gpu:%d' % 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)) # prepare the model with tf.device(gpu): # 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.cvpr2018_net(vol_size, nf_enc, nf_dec) # load initial weights if load_model_file is not None and load_model_file != '': print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # data generator # nb_gpus = len(gpu_id.split(',')) # assert np.mod(batch_size, nb_gpus) == 0, \ # 'batch_size should be a multiple of the nr. of gpus. ' + \ # 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus) nb_gpus = 1 train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size) atlas_vol_bs = np.repeat(atlas_vol, batch_size, axis=0) cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(gpu): # multi-gpu support if nb_gpus > 1: save_callback = nrn_gen.ModelCheckpointParallel(save_file_name) mg_model = multi_gpu_model(model, gpus=nb_gpus) # single-gpu else: save_callback = ModelCheckpoint(save_file_name) mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[data_loss, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit mg_model.fit_generator(cvpr2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def train_unsupervised_segmentation(data_dir, atlas_file, mapping_file, model, model_dir, gpu_id, lr, nb_epochs, init_stats, reg_param, steps_per_epoch, batch_size, stat_post_warp, warp_method, load_model_file, initial_epoch=0): """ model training function :param data_dir: folder with npz files (coregistered, intensity normalized) :param atlas_file: file with probabilistic atlas (coregistered to images) :param mapping_file: file with mapping from labels to tissue types :param model: registration (voxelmorph) model: vm1, vm2, or vm2double :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of epochs :param init_stats: file with guesses for means and log-variances (vectors init_mu, init_sigma) :param reg_param: smoothness/reconstruction tradeoff parameter (lambda in the paper) :param steps_per_epoch: frequency with which to save models :param batch_size: default of 1. can be larger, depends on GPU memory and volume size :param stat_post_warp: set to 1 to use warped atlas to estimate Gaussian parameters :param warp_method: set to 'WARP' if you want to warp the atlas :param load_model_file: optional h5 model file to initialize with :param initial_epoch: initial epoch """ # load reference soft edge and corresponding mask from provided files # (we used size 160x192x224). # Also: group labels in tissue types, if necessary if mapping_file is None: atlas_vol = np.load(atlas_file)['vol_data'][np.newaxis, ...] nb_labels = atlas_vol.shape[-1] else: atlas_full = np.load(atlas_file)['vol_data'][np.newaxis, ...] mapping = np.load(mapping_file)['mapping'].astype('int').flatten() assert len(mapping) == atlas_full.shape[-1], \ 'mapping shape %d is inconsistent with atlas shape %d' % (len(mapping), atlas_full.shape[-1]) nb_labels = 1 + np.max(mapping) atlas_vol = np.zeros( [1, *atlas_full.shape[1:-1], nb_labels.astype('int')]) for j in range(np.max(mapping.shape)): atlas_vol[0, ..., mapping[j]] = atlas_vol[0, ..., mapping[j]] + atlas_full[0, ..., j] vol_size = atlas_vol.shape[1:-1] # load guesses for means and variances init_mu = np.load(init_stats)['init_mu'] init_sigma = np.load(init_stats)['init_std'] # prepare data files train_vol_names = glob.glob(os.path.join(data_dir, '*.npz')) random.shuffle(train_vol_names) assert len(train_vol_names) > 0, "Could not find any training data" # 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] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] # prepare model and log folders if not os.path.isdir(model_dir): os.mkdir(model_dir) log_dir = os.path.join(model_dir, 'logs') if not os.path.isdir(log_dir): os.mkdir(log_dir) # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # prepare the model model = networks.cvpr2018_net_probatlas(vol_size, nf_enc, nf_dec, nb_labels, diffeomorphic=True, full_size=False, stat_post_warp=stat_post_warp, warp_method=warp_method, init_mu=init_mu, init_sigma=init_sigma) # load initial weights if load_model_file is not None: print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # data generator nb_gpus = len(gpu_id.split(',')) assert np.mod(batch_size, nb_gpus) == 0, \ 'batch_size should be a multiple of the nr. of gpus. ' + \ 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus) train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size) atlas_vol_bs = np.repeat(atlas_vol, batch_size, axis=0) cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(gpu): # multi-gpu support if nb_gpus > 1: save_callback = nrn_gen.ModelCheckpointParallel(save_file_name) mg_model = multi_gpu_model(model, gpus=nb_gpus) # single-gpu else: save_callback = ModelCheckpoint(save_file_name) mg_model = model # tensorBoard callback tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=True, write_images=False) # compile loss and parameters def data_loss(_, yp): m = tf.cast(model.inputs[0] > 0, tf.float32) return -K.sum(yp * m) / K.sum(m) if warp_method != 'WARP': reg_param = 0 # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[data_loss, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit mg_model.fit_generator(cvpr2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback, tensorboard], steps_per_epoch=steps_per_epoch, verbose=1)