drop_last=True)
valid_loader = data.DataLoader(valid_dataset,
batch_size=opt.batch_size,
num_workers=opt.workers,
pin_memory=opt.use_shared_memory)
#
# Instantiate the train loop
#
optimizer = setup_optimizer(model, opt.optimizer, opt.learning_rate)
train_loop_handler = TrainLoop(model,
train_loader,
valid_loader,
optimizer,
opt.backend,
gradient_clip=opt.gradient_clip,
use_tensorboard=opt.use_tensorboard)
#
# Add callbacks
#
train_loop_handler.add_callback([
IlluminationPredictorCallback(10, hdr_image_handler, ldr_image_handler,
train_dataset, opt, model.ae)
])
#
# Train the model
#
# Instantiate loaders
#
hdr_image_handler = HDRImageHandler(opt.mean_std,
perform_scale_perturbation=False)
test_dataset = AutoEncoderDataset(
opt.data_path, transform=hdr_image_handler.normalization_ops)
test_loader = data.DataLoader(test_dataset,
batch_size=1,
num_workers=0,
pin_memory=True)
#
# Instantiate the train loop
#
train_loop_handler = TrainLoop(model, None, test_loader, None, opt.backend)
train_loop_handler.setup_checkpoint(opt.load_best, opt.load_last,
opt.output_path, False)
#
# Add callbacks
#
train_loop_handler.add_callback(
[AutoEncoderCallback(10, hdr_image_handler, test_dataset, opt)])
#
# Test the model
#
train_loop_handler.test()
# Setup model
model = model_class(image_size=int(train_dataset.metadata["image_size"]))
if finetune_path != "None":
finetune_path = os.path.expandvars(finetune_path)
print("Finetuning path : {}".format(finetune_path))
checkpoint = torch.load(finetune_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
# Instantiate the train loop and train the model.
train_loop_handler = TrainLoop(model,
train_loader,
val_loader,
optimizer,
backend,
gradient_clip,
use_tensorboard=use_tensorboard,
tensorboard_log_path=tensorboard_path)
train_loop_handler.add_callback(callbacks)
print("Training Begins:")
train_loop_handler.loop(epochs,
output_path,
load_best_checkpoint=start_from_last,
save_all_checkpoints=False)
print("Training Complete")
configs=opt,
transform=ldr_image_handler,
dataset_purpose='test')
test_loader = data.DataLoader(test_dataset,
batch_size=1,
num_workers=opt.workers,
pin_memory=opt.use_shared_memory)
#
# Instantiate the train loop
#
train_loop_handler = TrainLoop(model,
None,
test_loader,
None,
opt.backend,
gradient_clip=False,
use_tensorboard=False)
train_loop_handler.setup_checkpoint(opt.load_best, opt.load_last,
opt.model_path, False)
#
# Add callbacks
#
train_loop_handler.add_callback([
IlluminationPredictorCallback(1,
hdr_image_handler,
ldr_image_handler,
test_dataset,
opt,