import torch, os, numpy as np, pandas as pd
from pathflowai.utils import *
#from large_data_utils import *
from pathflowai.datasets import *
from pathflowai.models import *
from pathflowai.schedulers import *
from pathflowai.visualize import *
import copy
from pathflowai.sampler import ImbalancedDatasetSampler
import argparse
import sqlite3
#from nonechucks import SafeDataLoader as DataLoader
from torch.utils.data import DataLoader
import click
import pysnooper
CONTEXT_SETTINGS = dict(help_option_names=['-h','--help'], max_content_width=90)
@click.group(context_settings= CONTEXT_SETTINGS)
@click.version_option(version='0.1')
def train():
pass
def return_model(training_opts):
model=generate_model(pretrain=training_opts['pretrain'],architecture=training_opts['architecture'],num_classes=training_opts['num_targets'], add_sigmoid=False, n_hidden=training_opts['n_hidden'], segmentation=training_opts['segmentation'])
if os.path.exists(training_opts['pretrained_save_location']) and not training_opts['predict']:
model_dict = torch.load(training_opts['pretrained_save_location'])
keys=list(model_dict.keys())
if not training_opts['segmentation']:
model_dict.update(dict(list(model.state_dict().items())[-2:]))#={k:model_dict[k] for k in keys[:-2]}
model.load_state_dict(model_dict) # this will likely break after pretraining?
elif os.path.exists(training_opts['save_location']) and training_opts['predict']:
model_dict = torch.load(training_opts['save_location'])
model.load_state_dict(model_dict)
if training_opts['extract_embedding']:
assert training_opts['extract_embedding']==training_opts['predict']
architecture=training_opts['architecture']
if hasattr(model,"fc"):
model.fc = model.fc[0]
elif hasattr(model,"output"):
model.output = model.output[0]
elif architecture.startswith('alexnet') or architecture.startswith('vgg') or architecture.startswith('densenet'):
model.classifier[6]=model.classifier[6][0]
if torch.cuda.is_available():
model.cuda()
return model
def run_test(dataloader):
for i,(X,y) in enumerate(dataloader):
np.save('X_test_{}.npy'.format(i),X.detach().cpu().numpy())#np.save('test_predictions.npy',model(X.cuda() if torch.cuda.is_available() else X).detach().cpu().numpy())
np.save('y_test_{}.npy'.format(i),y.detach().cpu().numpy())
if i==5:
exit()
def return_trainer_opts(model,training_opts,dataloaders,num_train_batches):
return dict(model=model,
n_epoch=training_opts['n_epoch'],
validation_dataloader=dataloaders['val'],
optimizer_opts=dict(name=training_opts['optimizer'],
lr=training_opts['lr'],
weight_decay=training_opts['wd']),
scheduler_opts=dict(scheduler=training_opts['scheduler_type'],
lr_scheduler_decay=0.5,
T_max=training_opts['T_max'],
eta_min=training_opts['eta_min'],
T_mult=training_opts['T_mult']),
loss_fn=training_opts['loss_fn'],
num_train_batches=num_train_batches,
seg_out_class=training_opts['seg_out_class'],
apex_opt_level=training_opts['apex_opt_level'],
checkpointing=training_opts['checkpointing'])
def return_transformer(training_opts):
dataset_df = pd.read_csv(training_opts['dataset_df']) if os.path.exists(training_opts['dataset_df']) else create_train_val_test(training_opts['train_val_test_splits'],training_opts['patch_info_file'],training_opts['patch_size'])
dataset_opts=dict(dataset_df=dataset_df, set='pass', patch_info_file=training_opts['patch_info_file'], input_dir=training_opts['input_dir'], target_names=training_opts['target_names'], pos_annotation_class=training_opts['pos_annotation_class'], segmentation=training_opts['segmentation'], patch_size=training_opts['patch_size'], fix_names=training_opts['fix_names'], other_annotations=training_opts['other_annotations'], target_segmentation_class=training_opts['target_segmentation_class'][0] if set=='train' else -1, target_threshold=training_opts['target_threshold'][0], oversampling_factor=training_opts['oversampling_factor'][0] if set=='train' else 1, n_segmentation_classes=training_opts['num_targets'],gdl=training_opts['loss_fn']=='gdl',mt_bce=training_opts['mt_bce'], classify_annotations=training_opts['classify_annotations'])
norm_dict = get_normalizer(training_opts['normalization_file'], dataset_opts)
transform_opts=dict(patch_size = training_opts['patch_resize'], mean=norm_dict['mean'], std=norm_dict['std'], resize=True, transform_platform=training_opts['transform_platform'] if not training_opts['segmentation'] else 'albumentations', user_transforms=training_opts['user_transforms'])
transformers = get_data_transforms(**transform_opts)
return dataset_df,dataset_opts,transformers
def return_datasets(training_opts,dataset_df,transformers):
datasets= {set: DynamicImageDataset(dataset_df, set, training_opts['patch_info_file'], transformers, training_opts['input_dir'], training_opts['target_names'], training_opts['pos_annotation_class'], segmentation=training_opts['segmentation'], patch_size=training_opts['patch_size'], fix_names=training_opts['fix_names'], other_annotations=training_opts['other_annotations'], target_segmentation_class=training_opts['target_segmentation_class'][0] if set=='train' else -1, target_threshold=training_opts['target_threshold'][0], oversampling_factor=training_opts['oversampling_factor'][0] if set=='train' else 1, n_segmentation_classes=training_opts['num_targets'],gdl=training_opts['loss_fn']=='gdl',mt_bce=training_opts['mt_bce'], classify_annotations=training_opts['classify_annotations'],dilation_jitter=training_opts['dilation_jitter'] if set == 'train' else {}) for set in ['train','val','test']}
# nc.SafeDataset(
print(datasets['train'])
if len(training_opts['target_segmentation_class']) > 1:
from functools import reduce
for i in range(1,len(training_opts['target_segmentation_class'])):
#print(training_opts['classify_annotations'])
datasets['train'].concat(DynamicImageDataset(dataset_df, 'train', training_opts['patch_info_file'], transformers, training_opts['input_dir'], training_opts['target_names'], training_opts['pos_annotation_class'], segmentation=training_opts['segmentation'], patch_size=training_opts['patch_size'], fix_names=training_opts['fix_names'], other_annotations=training_opts['other_annotations'], target_segmentation_class=training_opts['target_segmentation_class'][i], target_threshold=training_opts['target_threshold'][i], oversampling_factor=training_opts['oversampling_factor'][i],n_segmentation_classes=training_opts['num_targets'],gdl=training_opts['loss_fn']=='gdl',mt_bce=training_opts['mt_bce'],classify_annotations=training_opts['classify_annotations'],dilation_jitter=training_opts['dilation_jitter']))
#datasets['train']=reduce(lambda x,y: x.concat(y),[DynamicImageDataset(dataset_df, 'train', training_opts['patch_info_file'], transformers, training_opts['input_dir'], training_opts['target_names'], training_opts['pos_annotation_class'], segmentation=training_opts['segmentation'], patch_size=training_opts['patch_size'], fix_names=training_opts['fix_names'], other_annotations=training_opts['other_annotations'], target_segmentation_class=training_opts['target_segmentation_class'][i], target_threshold=training_opts['target_threshold'][i], oversampling_factor=training_opts['oversampling_factor'][i]) for i in range(len(training_opts['target_segmentation_class']))])
print(datasets['train'])
if training_opts['supplement']:
old_train_set = copy.deepcopy(datasets['train'])
datasets['train']=DynamicImageDataset(dataset_df, 'train', training_opts['patch_info_file'], transformers, training_opts['input_dir'], training_opts['target_names'], training_opts['pos_annotation_class'], segmentation=training_opts['segmentation'], patch_size=training_opts['patch_size'], fix_names=training_opts['fix_names'], other_annotations=training_opts['other_annotations'], target_segmentation_class=-1, target_threshold=training_opts['target_threshold'], oversampling_factor=1,n_segmentation_classes=training_opts['num_targets'],gdl=training_opts['loss_fn']=='gdl',mt_bce=training_opts['mt_bce'],classify_annotations=training_opts['classify_annotations'],dilation_jitter=training_opts['dilation_jitter'])
datasets['train'].concat(old_train_set)
if training_opts['subsample_p']<1.0:
datasets['train'].subsample(training_opts['subsample_p'])
if training_opts['subsample_p_val']<1.0:
if training_opts['subsample_p_val']==-1.:
training_opts['subsample_p_val']=training_opts['subsample_p']
if training_opts['subsample_p_val']<1.0:
datasets['val'].subsample(training_opts['subsample_p_val'])
if training_opts['classify_annotations']:
binarizer=datasets['train'].binarize_annotations(num_targets=training_opts['num_targets'],binary_threshold=training_opts['binary_threshold'])
datasets['val'].binarize_annotations(num_targets=training_opts['num_targets'],binary_threshold=training_opts['binary_threshold'])
datasets['test'].binarize_annotations(num_targets=training_opts['num_targets'],binary_threshold=training_opts['binary_threshold'])
training_opts['num_targets']=len(datasets['train'].targets)
for Set in ['train','val','test']:
print(datasets[Set].patch_info.iloc[:,6:].sum(axis=0))
if training_opts['prediction_set']!='test':
datasets['test']=datasets[training_opts['prediction_set']]
if training_opts['external_test_db'] and training_opts['external_test_dir']:
datasets['test'].update_dataset(input_dir=training_opts['external_test_dir'],new_db=training_opts['external_test_db'],prediction_basename=training_opts['prediction_basename'])
return datasets,training_opts,transform_opts
#@pysnooper.snoop('train_model.log')
def train_model_(training_opts):
"""Function to train, predict on model.
Parameters
----------
training_opts : dict
Training options populated from command line.
"""
cuda_available=torch.cuda.is_available()
model = return_model(training_opts)
dataset_df,dataset_opts,transformers=return_transformer(training_opts)
if training_opts['extract_embedding'] and training_opts['npy_file']:
dataset=NPYDataset(training_opts['patch_info_file'],training_opts['patch_size'],training_opts['npy_file'],transformers["test"])
dataset.embed(model,training_opts['batch_size'],training_opts['prediction_output_dir'])
exit()
datasets,training_opts,transform_opts=return_datasets(training_opts,dataset_df,transformers)
if training_opts['num_training_images_epoch']>0:
num_train_batches = min(training_opts['num_training_images_epoch'],len(datasets['train']))//training_opts['batch_size']
else:
num_train_batches = None
dataloaders={set: DataLoader(datasets[set], batch_size=training_opts['batch_size'], shuffle=(set=='train') if not (training_opts['imbalanced_correction'] and not training_opts['segmentation']) else False, num_workers=10, sampler=ImbalancedDatasetSampler(datasets[set]) if (training_opts['imbalanced_correction'] and set=='train' and not training_opts['segmentation']) else None) for set in ['train', 'val', 'test']}
print(dataloaders['train'].sampler) # FIXME VAL SEEMS TO BE MISSING DURING PREDICTION
print(dataloaders['val'].sampler)
if training_opts['run_test']: run_test(dataloaders['train'])
model_trainer_opts=return_trainer_opts(model,training_opts,dataloders,num_train_batches)
if not training_opts['predict']:
trainer = ModelTrainer(**model_trainer_opts)
if training_opts['imbalanced_correction2']:
trainer.add_class_balance_loss(datasets['train'])
elif training_opts['custom_weights']:
trainer.add_class_balance_loss(datasets['train'],custom_weights=training_opts['custom_weights'])
if training_opts['adopt_training_loss']:
trainer.val_loss_fn = trainer.loss_fn
trainer.fit(dataloaders['train'], verbose=True, print_every=1, plot_training_curves=True, plot_save_file=training_opts['training_curve'], print_val_confusion=training_opts['print_val_confusion'], save_val_predictions=training_opts['save_val_predictions'])
torch.save(trainer.model.state_dict(),training_opts['save_location'])
else:
if training_opts['extract_model']:
dataset_opts.update(dict(target_segmentation_class=-1, target_threshold=training_opts['target_threshold'][0] if len(training_opts['target_threshold']) else 0., set='test', binary_threshold=training_opts['binary_threshold'], num_targets=training_opts['num_targets'], oversampling_factor=1))
torch.save(dict(model=model,dataset_opts=dataset_opts, transform_opts=transform_opts),'{}.{}'.format(training_opts['save_location'],'extracted_model.pkl'))
exit()
trainer = ModelTrainer(**model_trainer_opts)
if training_opts['segmentation']:
for ID, dataset in (datasets['test'].split_by_ID() if not training_opts['prediction_basename'] else datasets['test'].select_IDs(training_opts['prediction_basename'])):
dataloader = DataLoader(dataset, batch_size=training_opts['batch_size'], shuffle=False, num_workers=10)
if training_opts['run_test']:
for X,y in dataloader:
np.save('test_predictions.npy',model(X.cuda() if torch.cuda.is_available() else X).detach().cpu().numpy())
exit()
y_pred = trainer.predict(dataloader)
print(ID,y_pred.shape)
segmentation_predictions2npy(y_pred, dataset.patch_info, dataset.segmentation_maps[ID], npy_output='{}/{}_predict.npy'.format(training_opts['prediction_output_dir'],ID), original_patch_size=training_opts['patch_size'], resized_patch_size=training_opts['patch_resize'], output_probs=(training_opts['seg_out_class']>=0))
else:
extract_embedding=training_opts['extract_embedding']
if extract_embedding:
trainer.bce=False
y_pred = trainer.predict(dataloaders['test'])
patch_info = dataloaders['test'].dataset.patch_info
if extract_embedding:
patch_info['name']=patch_info.astype(str).apply(lambda x: '\n'.join(['{}:{}'.format(k,v) for k,v in x.to_dict().items()]),axis=1)#.apply(','.join,axis=1)
embeddings=pd.DataFrame(y_pred,index=patch_info['name'])
embeddings['ID']=patch_info['ID'].values
torch.save(dict(embeddings=embeddings,patch_info=patch_info),join(training_opts['prediction_output_dir'],'embeddings.pkl'))
else:
if len(y_pred.shape)>1 and y_pred.shape[1]>1:
annotations = np.vectorize(lambda x: x+'_pred')(np.arange(y_pred.shape[1]).astype(str)).tolist() # [training_opts['pos_annotation_class']]+training_opts['other_annotations']] if training_opts['classify_annotations'] else
for i in range(y_pred.shape[1]):
patch_info.loc[:,annotations[i]]=y_pred[:,i]
patch_info['y_pred']=y_pred if (training_opts['num_targets']==1 or not (training_opts['classify_annotations'] or training_opts['mt_bce'])) else y_pred.argmax(axis=1)
conn = sqlite3.connect(training_opts['prediction_save_path'])
patch_info.to_sql(str(training_opts['patch_size']),con=conn, if_exists=('replace')) # if not training_opts['prediction_basename'] else 'append'))
conn.close()
@train.command()
@click.option('-s', '--segmentation', is_flag=True, help='Segmentation task.', show_default=True)
@click.option('-p', '--prediction', is_flag=True, help='Predict on model.', show_default=True)
@click.option('-pa', '--pos_annotation_class', default='', help='Annotation Class from which to apply positive labels.', type=click.Path(exists=False), show_default=True)
@click.option('-oa', '--other_annotations', default=[], multiple=True, help='Annotations in image.', type=click.Path(exists=False), show_default=True)
@click.option('-o', '--save_location', default='', help='Model Save Location, append with pickle .pkl.', type=click.Path(exists=False), show_default=True)
@click.option('-pt', '--pretrained_save_location', default='', help='Model Save Location, append with pickle .pkl, pretrained by previous analysis to be finetuned.', type=click.Path(exists=False), show_default=True)
@click.option('-i', '--input_dir', default='', help='Input directory containing slides and everything.', type=click.Path(exists=False), show_default=True)
@click.option('-ps', '--patch_size', default=224, help='Patch size.', show_default=True)
@click.option('-pr', '--patch_resize', default=224, help='Patch resized.', show_default=True)
@click.option('-tg', '--target_names', default=[], multiple=True, help='Targets.', type=click.Path(exists=False), show_default=True)
@click.option('-df', '--dataset_df', default='', help='CSV file with train/val/test and target info.', type=click.Path(exists=False), show_default=True)
@click.option('-fn', '--fix_names', is_flag=True, help='Whether to fix names in dataset_df.', show_default=True)
@click.option('-a', '--architecture', default='alexnet', help='Neural Network Architecture.', type=click.Choice(['alexnet', 'densenet121', 'densenet161', 'densenet169', 'densenet201',
'inception_v3', 'resnet101', 'resnet152', 'resnet18', 'resnet34', 'resnet50', 'vgg11', 'vgg11_bn','unet','unet2','nested_unet','fast_scnn',
'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn', 'deeplabv3_resnet101','deeplabv3_resnet50','fcn_resnet101', 'fcn_resnet50',"sqnxt23_w3d2", "sqnxt23_w2", "sqnxt23v5_w1", "sqnxt23v5_w3d2", "sqnxt23v5_w2"]+['efficientnet-b{}'.format(i) for i in range(8)]), show_default=True)
@click.option('-imb', '--imbalanced_correction', is_flag=True, help='Attempt to correct for imbalanced data.', show_default=True)
@click.option('-imb2', '--imbalanced_correction2', is_flag=True, help='Attempt to correct for imbalanced data.', show_default=True)
@click.option('-ca', '--classify_annotations', is_flag=True, help='Classify annotations.', show_default=True)
@click.option('-nt', '--num_targets', default=1, help='Number of targets.', show_default=True)
@click.option('-ss', '--subsample_p', default=1.0, help='Subsample training set.', show_default=True)
@click.option('-ssv', '--subsample_p_val', default=-1., help='Subsample val set. If not set, defaults to that of training set', show_default=True)
@click.option('-t', '--num_training_images_epoch', default=-1, help='Number of training images per epoch. -1 means use all training images each epoch.', show_default=True)
@click.option('-lr', '--learning_rate', default=1e-2, help='Learning rate.', show_default=True)
@click.option('-tp', '--transform_platform', default='torch', help='Transform platform for nonsegmentation tasks.', type=click.Choice(['torch','albumentations']))
@click.option('-ne', '--n_epoch', default=10, help='Number of epochs.', show_default=True)
@click.option('-pi', '--patch_info_file', default='patch_info.db', help='Patch info file.', type=click.Path(exists=False), show_default=True)
@click.option('-tc', '--target_segmentation_class', default=[-1], multiple=True, help='Segmentation Class to finetune on.', show_default=True)
@click.option('-tt', '--target_threshold', default=[0.], multiple=True, help='Threshold to include target for segmentation if saving one class.', show_default=True)
@click.option('-ov', '--oversampling_factor', default=[1.], multiple=True, help='How much to oversample training set.', show_default=True)
@click.option('-sup', '--supplement', is_flag=True, help='Use the thresholding to supplement the original training set.', show_default=True)
@click.option('-bs', '--batch_size', default=10, help='Batch size.', show_default=True)
@click.option('-rt', '--run_test', is_flag=True, help='Output predictions for a batch to "test_predictions.npy". Use for debugging.', show_default=True)
@click.option('-mtb', '--mt_bce', is_flag=True, help='Run multi-target bce predictions on the annotations.', show_default=True)
@click.option('-po', '--prediction_output_dir', default='predictions', help='Where to output segmentation predictions.', type=click.Path(exists=False), show_default=True)
@click.option('-ee', '--extract_embedding', is_flag=True, help='Extract embeddings.', show_default=True)
@click.option('-em', '--extract_model', is_flag=True, help='Save entire torch model.', show_default=True)
@click.option('-bt', '--binary_threshold', default=0., help='If running binary classification on annotations, dichotomize selected annotation as such.', show_default=True)
@click.option('-prt', '--pretrain', is_flag=True, help='Pretrain on ImageNet.', show_default=True)
@click.option('-olf', '--overwrite_loss_fn', default='', help='Overwrite the default training loss functions with loss of choice.', type=click.Choice(['','bce','mse','focal','dice','gdl','ce']), show_default=True)
@click.option('-atl', '--adopt_training_loss', is_flag=True, help='Adopt training loss function for validation calculation.', show_default=True)
@click.option('-tdb', '--external_test_db', default='', help='External database of samples to test on.', type=click.Path(exists=False), show_default=True)
@click.option('-tdir', '--external_test_dir', default='', help='External directory of samples to test on.', type=click.Path(exists=False), show_default=True)
@click.option('-pb', '--prediction_basename', default=[''], multiple=True, help='For segmentation tasks, if supplied, can predict on these basenames rather than the entire test set. Only works for segmentation tasks for now', show_default=True)
@click.option('-cw', '--custom_weights', default='', help='Comma delimited custom weights', type=click.Path(exists=False), show_default=True)
@click.option('-pset', '--prediction_set', default='test', help='Dataset to predict on.', type=click.Choice(['train','val','test']), show_default=True)
@click.option('-ut', '--user_transforms_file', default='', help='YAML file to add transforms from.', type=click.Path(exists=False), show_default=True)
@click.option('-svp', '--save_val_predictions', is_flag=True, help='Whether to save the validation predictions.', show_default=True)
@click.option('-soc', '--seg_out_class', default=-1, help='Output a particular segmentation class probabilities.', show_default=True)
@click.option('-aol', '--apex_opt_level', default='O2', help='YAML file to add transforms from.', type=click.Choice(['O0','O1','O2','O3']), show_default=True)
@click.option('-ckp', '--checkpointing', is_flag=True, help='Save intermediate models to ./checkpoints.', show_default=True)
@click.option('-npy', '--npy_file', default='', help='Specify one file to extract embeddings from. Embeddings are output into predictions directory', type=click.Path(exists=False), show_default=True)
@click.option('-gpu', '--gpu_id', default=-1, help='Set GPU if 0 and greater.', show_default=True)
def train_model(segmentation,prediction,pos_annotation_class,other_annotations,save_location,pretrained_save_location,input_dir,patch_size,patch_resize,target_names,dataset_df,fix_names, architecture, imbalanced_correction, imbalanced_correction2, classify_annotations, num_targets, subsample_p,subsample_p_val,num_training_images_epoch, learning_rate, transform_platform, n_epoch, patch_info_file, target_segmentation_class, target_threshold, oversampling_factor, supplement, batch_size, run_test, mt_bce, prediction_output_dir, extract_embedding, extract_model, binary_threshold, pretrain, overwrite_loss_fn, adopt_training_loss, external_test_db,external_test_dir, prediction_basename, custom_weights, prediction_set, user_transforms_file, save_val_predictions, seg_out_class, apex_opt_level, checkpointing, npy_file, gpu_id):
"""Train and predict using model for regression and classification tasks."""
# add separate pretrain ability on separating cell types, then transfer learn
# add pretrain and efficient net, pretraining remove last layer while loading state dict
if gpu_id>=0:
torch.cuda.set_device(gpu_id)
target_segmentation_class=list(map(int,target_segmentation_class))
target_threshold=list(map(float,target_threshold))
oversampling_factor=[(int(x) if float(x)>=1 else float(x)) for x in oversampling_factor]
other_annotations=list(other_annotations)
prediction_basename=list(filter(None,prediction_basename))
command_opts = dict(segmentation=segmentation,
prediction=prediction,
pos_annotation_class=pos_annotation_class,
other_annotations=other_annotations,
save_location=save_location,
pretrained_save_location=pretrained_save_location,
input_dir=input_dir,
patch_size=patch_size,
target_names=target_names,
dataset_df=dataset_df,
fix_names=fix_names,
architecture=architecture,
patch_resize=patch_resize,
imbalanced_correction=imbalanced_correction,
imbalanced_correction2=imbalanced_correction2,
classify_annotations=classify_annotations,
num_targets=num_targets,
subsample_p=subsample_p,
num_training_images_epoch=num_training_images_epoch,
lr=learning_rate,
transform_platform=transform_platform,
n_epoch=n_epoch,
patch_info_file=patch_info_file,
target_segmentation_class=target_segmentation_class,
target_threshold=target_threshold,
oversampling_factor=oversampling_factor,
supplement=supplement,
predict=prediction,
batch_size=batch_size,
run_test=run_test,
mt_bce=mt_bce,
prediction_output_dir=prediction_output_dir,
extract_embedding=extract_embedding,
extract_model=extract_model,
binary_threshold=binary_threshold,
subsample_p_val=subsample_p_val,
wd=1e-3,
scheduler_type='warm_restarts',
T_max=10,
T_mult=2,
eta_min=5e-8,
optimizer='adam',
n_hidden=100,
pretrain=pretrain,
training_curve='training_curve.png',
adopt_training_loss=adopt_training_loss,
external_test_db=external_test_db,
external_test_dir=external_test_dir,
prediction_basename=prediction_basename,
save_val_predictions=save_val_predictions,
custom_weights=custom_weights,
prediction_set=prediction_set,
user_transforms=dict(),
dilation_jitter=dict(),
seg_out_class=seg_out_class,
apex_opt_level=apex_opt_level,
checkpointing=checkpointing,
npy_file=npy_file)
training_opts = dict(normalization_file="normalization_parameters.pkl",
loss_fn='bce',
print_val_confusion=True,
prediction_save_path = 'predictions.db',
train_val_test_splits='train_val_test.pkl'
)
segmentation_training_opts = copy.deepcopy(training_opts)
segmentation_training_opts.update(dict(loss_fn='dice',#gdl dice+ce
normalization_file='normalization_segmentation.pkl',
fix_names=False,
))
if segmentation:
training_opts = segmentation_training_opts
for k in command_opts:
training_opts[k] = command_opts[k]
if classify_annotations:
if training_opts['num_targets']==1:
training_opts['loss_fn']='bce'
else:
training_opts['loss_fn']='ce'
if mt_bce:
training_opts['loss_fn']='bce'
if overwrite_loss_fn:
training_opts['loss_fn']=overwrite_loss_fn
if user_transforms_file and os.path.exists(user_transforms_file):
from yaml import load as yml_load
try:
from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
from yaml import Loader, Dumper
with open(user_transforms_file) as f:
training_opts['user_transforms']=yml_load(f,Loader=Loader)
if 'dilationjitter' in list(training_opts['user_transforms'].keys()):
training_opts['dilation_jitter']=training_opts['user_transforms'].pop('dilationjitter')
train_model_(training_opts)
if __name__=='__main__':
train()
Datasets
Models
