def get_data_loaders(args):
train_dataset = UnpairedDataset(args.data_dir,
phase='train',
shuffle_pairs=True,
resize_shape=args.resize_shape,
crop_shape=args.crop_shape,
direction=args.direction)
train_loader = DataLoader(train_dataset,
args.batch_size,
shuffle=True,
num_workers=args.num_workers)
val_dataset = PairedDataset(args.data_dir,
phase='val',
resize_shape=args.resize_shape,
crop_shape=args.crop_shape,
direction=args.direction)
val_loader = DataLoader(val_dataset,
args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=True)
return train_loader, val_loader
if __name__ == '__main__':
parser = TrainArgParser()
train(parser.parse_args())
if val_macro_dice.result() > best_val_dice:
best_val_dice = val_macro_dice.result()
patience = 0
if args.save_folder:
model.save_weights(os.path.join(args.save_folder,
'chkpt.hdf5'))
print('Saved model weights.', flush=True)
elif patience == args.patience:
print(
'Validation dice has not improved in {} epochs. Stopped training.'
.format(args.patience),
flush=True)
return
else:
patience += 1
# Reset statistics.
train_loss.reset_states()
train_macro_dice.reset_states()
train_micro_dice.reset_states()
val_loss.reset_states()
val_macro_dice.reset_states()
val_micro_dice.reset_states()
if __name__ == '__main__':
parser = TrainArgParser()
args = parser.parse_args()
print('Train args: {}'.format(args), flush=True)
train(args)
from args import TrainArgParser
import util
def train(args):
"""Run training loop with the given args.
The function consists of the following steps:
1. Load model: gets the model from a checkpoint or from models/models.py.
2. Load optimizer and learning rate scheduler.
3. Get data loaders and class weights.
4. Get loss functions: cross entropy loss and weighted loss functions.
5. Get logger, evaluator, and saver.
6. Run training loop, evaluate and save model periodically.
"""
model_args = args.model_args
logger_args = args.logger_args
optim_args = args.optim_args
data_args = args.data_args
transform_args = args.transform_args
print(args)
if __name__ == '__main__':
parser = TrainArgParser()
args = util.get_auto_args(parser)
train(args)
#train(parser.parse_args())