def main_train(model, stage):
# setting tensorboard
tensorboard_path = 'runs/{}_{}'.format(stage, args.saved_model_name)
tb_writer = TensorboardWriter(tensorboard_path)
# get one testing image, used to visualize result in each eopch
testing_rgb, testing_lidar, testing_mask = tb_writer.get_testing_img()
testing_rgb, testing_lidar, testing_mask = testing_rgb.to(
DEVICE), testing_lidar.to(DEVICE), testing_mask.to(DEVICE)
# setting early stop, if result doen't improve more than PATIENCE times, stop iteration
early_stop = EarlyStop(patience=10, mode='min')
# get data loader
loader = {'train': get_loader('train', num_data=args.num_data), \
'val': get_loader('val', shuffle=False, num_data=1000)}
for epoch in range(args.epoch):
saved_model_path = os.path.join(
SAVED_MODEL_PATH, "{}_{}_e{}".format(args.saved_model_name, stage,
epoch + 1))
train_losses, val_losses = train_val(model, loader, epoch, DEVICE,
stage)
# predict dense and surface normal using testing image and write them to tensorboard
predicted_dense, pred_surface_normal = get_depth_and_normal(
model, testing_rgb, testing_lidar, testing_mask)
tb_writer.tensorboard_write(epoch, train_losses, val_losses,
predicted_dense, pred_surface_normal)
if early_stop.stop(val_losses[0], model, epoch + 1, saved_model_path):
break
tb_writer.close()
def main(config):
if config.checkpoint_dir is None:
config.checkpoint_dir = 'checkpoints'
if not os.path.exists(
os.path.join(config.checkpoint_dir, config.style_image_name)):
os.makedirs(
os.path.join(config.checkpoint_dir, config.style_image_name))
if not os.path.exists(
os.path.join(config.sample_dir, config.style_image_name)):
os.makedirs(os.path.join(config.sample_dir, config.style_image_name))
if not os.path.exists(config.data_path) or not os.listdir(
config.data_path):
download_dataset(config)
print(f"photo to {config.style_dir} style transfer using cnn")
data_loader, val_data_loader = get_loader(config.data_path,
config.image_size,
config.batch_size,
config.sample_batch_size)
trainer = Trainer(config, data_loader)
trainer.train()
tester = Tester(config, val_data_loader)
tester.test()
def main(config):
cudnn.benchmark = True
data_loader = get_loader(config.model, config.image_path,
config.metadata_path, config.mode,
config.batch_size, config.shuffle)
solver = Solver(data_loader, config)
if config.mode == 'train':
solver.train()
elif config.mode == 'test':
solver.test()
def main():
# initialize tensorboard for visualization training progress
tensorboard_path = 'runs/{}'.format(args.saved_model_name)
tb_writer = TensorboardWriter(tensorboard_path)
# set the early stop
early_stop = EarlyStop(patience=5, mode='max')
# get train and val dataloader
train_loader = get_loader('train',
batch_size=32,
shuffle=True,
num_data=-1)
val_loader = get_loader('val', batch_size=32, shuffle=False, num_data=-1)
loader = dict(train=train_loader, val=val_loader)
# choose the classifier
clf = resnet18(pretrained=False, num_classes=1, in_channels=2).to(DEVICE)
#clf = vgg11_bn(pretrained=False, num_classes=1).to(DEVICE)
#clf = vgg11(pretrained=False, num_classes=1).to(DEVICE)
# set criterion and optimizer
criterion = nn.BCELoss(reduction='mean')
optimizer = optim.Adam(clf.parameters(), lr=1e-4)
# training
for epoch in range(1, 100):
train_loss, val_loss, train_acc, val_acc = train_val(
clf, criterion, optimizer, loader, epoch, DEVICE)
tb_writer.tensorboard_write(epoch, train_loss, val_loss, train_acc,
val_acc)
saved_model_path = os.path.join("saved_model",
"{}".format(args.saved_model_name))
if early_stop.stop(val_acc, clf, epoch, saved_model_path):
break
tb_writer.close()
def main(config):
if not os.path.exists(config.checkpoint_dir):
os.makedirs(config.checkpoint_dir)
print('[*] Load Dataset')
train_loader, test_loader = get_loader(config.dataset, config.image_size,
config.batch_size)
print('[*] Train')
trainer = Trainer(config, train_loader)
trainer.train()
print('[*] Test')
tester = Tester(config, test_loader)
tester.test()
def main(config):
# make directory not existed
if config.checkpoint_dir is None:
config.checkpoint_dir = 'checkpoints'
if not os.path.exists(config.checkpoint_dir):
os.makedirs(config.checkpoint_dir)
if not os.path.exists(config.sample_dir):
os.makedirs(config.sample_dir)
print(f"ESRGAN start")
data_loader = get_loader(config.image_size, config.scale_factor,
config.batch_size, config.sample_batch_size)
trainer = Trainer(config, data_loader)
trainer.train()
def main(config):
if config.checkpoint_dir is None:
config.checkpoint_dir = 'checkpoints'
if not os.path.exists(config.checkpoint_dir):
os.makedirs(config.checkpoint_dir)
if not os.path.exists(config.sample_dir):
os.makedirs(config.sample_dir)
print(f"{config.style} CycleGAN")
from_style = config.style.split('2')[0]
to_style = config.style.split('2')[1]
data_loader = get_loader(from_style, to_style, config)
trainer = Trainer(config, data_loader)
trainer.train()
def main(config):
if config.checkpoint_dir is None:
config.checkpoint_dir = 'checkpoints'
if not os.path.exists(config.checkpoint_dir):
os.makedirs(config.checkpoint_dir)
if not os.path.exists(config.sample_dir):
os.makedirs(config.sample_dir)
print(f"{config.from_style} to {config.to_style} CycleGAN")
data_loader, val_data_loader = get_loader(config.from_style,
config.to_style, config)
trainer = Trainer(config, data_loader)
trainer.train()
tester = Tester(config, val_data_loader)
tester.test()
def main(config):
if config.checkpoint_dir is None:
config.checkpoint_dir = 'checkpoints'
if not os.path.exists(config.checkpoint_dir):
os.makedirs(config.checkpoint_dir)
if not os.path.exists(config.sample_dir):
os.makedirs(config.sample_dir)
print(f"photo to {config.style_dir} style transfer using cnn")
data_loader, val_data_loader = get_loader(config.from_style,
config.to_style,
config.batch_size,
config.sample_batch_size)
trainer = Trainer(config, data_loader)
trainer.train()
tester = Tester(config, val_data_loader)
tester.test()
def get_testing_img(self):
"""Load one example for visualizing result in each epoch.
Use tensorboard to visualize this example
Returns:
rgb: rgb input of the model (1 x c x w x h)
lidar: lidar input of the model
mask: mask input of the model
"""
loader = get_loader('val', shuffle=False, num_data=1, crop=False)
for rgb, lidar, mask, gt_depth, gt_surface_normal in loader:
self.writer.add_image('RGB input', rgb[0] / 255.0, 1)
self.writer.add_image('lidar input', lidar[0], 1)
self.writer.add_image('GroundTruth depth',
normal_to_0_1(gt_depth[0]), 1)
self.gt_mask = torch.tensor(
np.where(gt_depth.numpy() > 0.0, 1.0, 0.0)) # b x 1 x w x h
#self.gt_normal_mask = gt_normal_mask # b x 1 x w x h
return rgb, lidar, mask, gt_surface_normal