def main():
global best_error, n_iter, device
args = parser.parse_args()
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = Path(args.name)
args.save_path = 'checkpoints' / save_path / timestamp
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
np.random.seed(args.seed)
cudnn.deterministic = True
cudnn.benchmark = True
training_writer = SummaryWriter(args.save_path)
output_writers = []
if args.log_output:
for i in range(3):
output_writers.append(
SummaryWriter(args.save_path / 'valid' / str(i)))
# Data loading code
normalize = custom_transforms.Normalize(mean=[0.45, 0.45, 0.45],
std=[0.225, 0.225, 0.225])
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = SequenceFolder(args.data,
transform=train_transform,
seed=args.seed,
train=True,
sequence_length=args.sequence_length)
# if no Groundtruth is avalaible, Validation set is the same type as training set to measure photometric loss from warping
if args.with_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(args.data, transform=valid_transform)
else:
val_set = SequenceFolder(args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length)
print('{} samples found in {} train scenes'.format(len(train_set),
len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
print("=> creating model")
disp_net = models.DispResNet(args.resnet_layers,
args.with_pretrain).to(device)
pose_net = models.PoseResNet(18, args.with_pretrain).to(device)
# load parameters
if args.pretrained_disp:
print("=> using pre-trained weights for DispResNet")
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'], strict=False)
if args.pretrained_pose:
print("=> using pre-trained weights for PoseResNet")
weights = torch.load(args.pretrained_pose)
pose_net.load_state_dict(weights['state_dict'], strict=False)
disp_net = torch.nn.DataParallel(disp_net)
pose_net = torch.nn.DataParallel(pose_net)
print('=> setting adam solver')
optim_params = [{
'params': disp_net.parameters(),
'lr': args.lr
}, {
'params': pose_net.parameters(),
'lr': args.lr
}]
optimizer = torch.optim.Adam(optim_params,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([
'train_loss', 'photo_loss', 'smooth_loss',
'geometry_consistency_loss'
])
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
for epoch in range(args.epochs):
logger.epoch_bar.update(epoch)
# train for one epoch
logger.reset_train_bar()
train_loss = train(args, train_loader, disp_net, pose_net, optimizer,
args.epoch_size, logger, training_writer)
logger.train_writer.write(' * Avg Loss : {:.3f}'.format(train_loss))
# evaluate on validation set
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
epoch, logger,
output_writers)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_net,
epoch, logger,
output_writers)
error_string = ', '.join('{} : {:.3f}'.format(name, error)
for name, error in zip(error_names, errors))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error, epoch)
# Up to you to chose the most relevant error to measure your model's performance, careful some measures are to maximize (such as a1,a2,a3)
decisive_error = errors[1]
if best_error < 0:
best_error = decisive_error
# remember lowest error and save checkpoint
is_best = decisive_error < best_error
best_error = min(best_error, decisive_error)
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_net.module.state_dict()
}, is_best)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])
logger.epoch_bar.finish()
def main():
global args, best_error, n_iter, device
args = parser.parse_args()
save_path = save_path_formatter(args, parser)
args.save_path = 'checkpoints_shifted' / save_path
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
training_writer = SummaryWriter(args.save_path)
output_writers = []
if args.log_output:
for i in range(3):
output_writers.append(
SummaryWriter(args.save_path / 'valid' / str(i)))
# Data loading code
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = ShiftedSequenceFolder(
args.data,
transform=train_transform,
seed=args.seed,
train=True,
sequence_length=args.sequence_length,
target_displacement=args.target_displacement)
# if no Groundtruth is avalaible, Validation set is the same type as training set to measure photometric loss from warping
if args.with_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(args.data, transform=valid_transform)
else:
val_set = SequenceFolder(
args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
)
print('{} samples found in {} train scenes'.format(len(train_set),
len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
adjust_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=0,
pin_memory=True
) # workers is set to 0 to avoid multiple instances to be modified at the same time
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
train.args = args
# create model
print("=> creating model")
disp_net = models.DispNetS().cuda()
output_exp = args.mask_loss_weight > 0
if not output_exp:
print("=> no mask loss, PoseExpnet will only output pose")
pose_exp_net = models.PoseExpNet(
nb_ref_imgs=args.sequence_length - 1,
output_exp=args.mask_loss_weight > 0).to(device)
if args.pretrained_exp_pose:
print("=> using pre-trained weights for explainabilty and pose net")
weights = torch.load(args.pretrained_exp_pose)
pose_exp_net.load_state_dict(weights['state_dict'], strict=False)
else:
pose_exp_net.init_weights()
if args.pretrained_disp:
print("=> using pre-trained weights for Dispnet")
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'])
else:
disp_net.init_weights()
cudnn.benchmark = True
disp_net = torch.nn.DataParallel(disp_net)
pose_exp_net = torch.nn.DataParallel(pose_exp_net)
print('=> setting adam solver')
parameters = chain(disp_net.parameters(), pose_exp_net.parameters())
optimizer = torch.optim.Adam(parameters,
args.lr,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(
['train_loss', 'photo_loss', 'explainability_loss', 'smooth_loss'])
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
for epoch in range(args.epochs):
logger.epoch_bar.update(epoch)
# train for one epoch
logger.reset_train_bar()
train_loss = train(args, train_loader, disp_net, pose_exp_net,
optimizer, args.epoch_size, logger, training_writer)
logger.train_writer.write(' * Avg Loss : {:.3f}'.format(train_loss))
if (epoch + 1) % 5 == 0:
train_set.adjust = True
logger.reset_train_bar(len(adjust_loader))
average_shifts = adjust_shifts(args, train_set, adjust_loader,
pose_exp_net, epoch, logger,
training_writer)
shifts_string = ' '.join(
['{:.3f}'.format(s) for s in average_shifts])
logger.train_writer.write(
' * adjusted shifts, average shifts are now : {}'.format(
shifts_string))
for i, shift in enumerate(average_shifts):
training_writer.add_scalar('shifts{}'.format(i), shift, epoch)
train_set.adjust = False
# evaluate on validation set
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
epoch, logger,
output_writers)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_exp_net,
epoch, logger,
output_writers)
error_string = ', '.join('{} : {:.3f}'.format(name, error)
for name, error in zip(error_names, errors))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error, epoch)
# Up to you to chose the most relevant error to measure your model's performance, careful some measures are to maximize (such as a1,a2,a3)
decisive_error = errors[0]
if best_error < 0:
best_error = decisive_error
# remember lowest error and save checkpoint
is_best = decisive_error < best_error
best_error = min(best_error, decisive_error)
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_exp_net.module.state_dict()
}, is_best)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])
logger.epoch_bar.finish()
def main():
global opt, best_prec1
opt = parser.parse_args()
print(opt)
# Data loading
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(),
custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
valid_transform = custom_transforms.Compose([
custom_transforms.ArrayToTensor(),
custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
print('Loading scenes in', opt.data_dir)
train_set = SequenceFolder(opt.data_dir,
transform=train_transform,
seed=opt.seed,
train=True,
sequence_length=opt.sequence_length)
val_set = ValidationSet(opt.data_dir, transform=valid_transform)
print(len(train_set), 'samples found')
print(len(val_set), 'samples found')
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.workers,
pin_memory=True)
# val_loader = torch.utils.data.DataLoader(val_set, batch_size=opt.batch_size,
# shuffle=False, num_workers=opt.workers,
# pin_memory=True)
if opt.epoch == 0:
opt.epoch_size = len(train_loader)
# Done loading
disp_model = dispnet.DispNet().cuda()
pose_model = posenet.PoseNet().cuda()
disp_model, pose_model, optimizer = init.setup(disp_model, pose_model, opt)
print(disp_model, pose_model)
trainer = train.Trainer(disp_model, pose_model, optimizer, opt)
if opt.resume:
if os.path.isfile(opt.resume):
# disp_model, pose_model, optimizer, opt, best_prec1 = init.resumer(opt, disp_model, pose_model, optimizer)
disp_model, pose_model, optimizer, opt = init.resumer(
opt, disp_model, pose_model, optimizer)
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
cudnn.benchmark = True
for epoch in range(opt.start_epoch, opt.epochs):
utils.adjust_learning_rate(opt, optimizer, epoch)
print("Starting epoch number:", epoch + 1, "Learning rate:",
optimizer.param_groups[0]["lr"])
if opt.testOnly == False:
trainer.train(train_loader, epoch, opt)
# init.save_checkpoint(opt, disp_model, pose_model, optimizer, best_prec1, epoch)
init.save_checkpoint(opt, disp_model, pose_model, optimizer, epoch)
def main():
global best_error, n_iter, device
args = parser.parse_args()
if args.dataset_format == 'stacked':
from datasets.stacked_sequence_folders import SequenceFolder
elif args.dataset_format == 'sequential':
from datasets.sequence_folders import SequenceFolder
save_path = save_path_formatter(args, parser)
args.save_path = 'checkpoints' / save_path
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
if args.evaluate:
args.epochs = 0
tb_writer = SummaryWriter(args.save_path)
# Data loading code
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = SequenceFolder(args.data,
transform=train_transform,
seed=args.seed,
train=True,
sequence_length=args.sequence_length)
# if no Groundtruth is avalaible, Validation set is the same type as training set to measure photometric loss from warping
if args.with_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(args.data, transform=valid_transform)
else:
val_set = SequenceFolder(
args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
)
print('{} samples found in {} train scenes'.format(len(train_set),
len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
print("=> creating model")
disp_net = models.DispNetS().to(device)
seg_net = DeepLab(num_classes=args.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn).to(device)
if args.pretrained_seg:
print("=> using pre-trained weights for seg net")
weights = torch.load(args.pretrained_seg)
seg_net.load_state_dict(weights, strict=False)
output_exp = args.mask_loss_weight > 0
if not output_exp:
print("=> no mask loss, PoseExpnet will only output pose")
pose_exp_net = models.PoseExpNet(
nb_ref_imgs=args.sequence_length - 1,
output_exp=args.mask_loss_weight > 0).to(device)
if args.pretrained_exp_pose:
print("=> using pre-trained weights for explainabilty and pose net")
weights = torch.load(args.pretrained_exp_pose)
pose_exp_net.load_state_dict(weights['state_dict'], strict=False)
else:
pose_exp_net.init_weights()
if args.pretrained_disp:
print("=> using pre-trained weights for Dispnet")
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'])
else:
disp_net.init_weights()
cudnn.benchmark = True
disp_net = torch.nn.DataParallel(disp_net)
pose_exp_net = torch.nn.DataParallel(pose_exp_net)
seg_net = torch.nn.DataParallel(seg_net)
print('=> setting adam solver')
optim_params = [{
'params': disp_net.parameters(),
'lr': args.lr
}, {
'params': pose_exp_net.parameters(),
'lr': args.lr
}]
optimizer = torch.optim.Adam(optim_params,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(
['train_loss', 'photo_loss', 'explainability_loss', 'smooth_loss'])
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
if args.pretrained_disp or args.evaluate:
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
0, logger, tb_writer)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_exp_net,
0, logger, tb_writer)
for error, name in zip(errors, error_names):
tb_writer.add_scalar(name, error, 0)
error_string = ', '.join(
'{} : {:.3f}'.format(name, error)
for name, error in zip(error_names[2:9], errors[2:9]))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for epoch in range(args.epochs):
logger.epoch_bar.update(epoch)
# train for one epoch
logger.reset_train_bar()
train_loss = train(args, train_loader, disp_net, pose_exp_net, seg_net,
optimizer, args.epoch_size, logger, tb_writer)
logger.train_writer.write(' * Avg Loss : {:.3f}'.format(train_loss))
# evaluate on validation set
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
seg_net, epoch, logger,
tb_writer)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_exp_net,
epoch, logger, tb_writer)
error_string = ', '.join('{} : {:.3f}'.format(name, error)
for name, error in zip(error_names, errors))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for error, name in zip(errors, error_names):
tb_writer.add_scalar(name, error, epoch)
# Up to you to chose the most relevant error to measure your model's performance, careful some measures are to maximize (such as a1,a2,a3)
decisive_error = errors[1]
if best_error < 0:
best_error = decisive_error
# remember lowest error and save checkpoint
is_best = decisive_error < best_error
best_error = min(best_error, decisive_error)
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_exp_net.module.state_dict()
}, is_best)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])
logger.epoch_bar.finish()
def main():
global best_error, n_iter, device
args = parser.parse_args()
if args.dataset_format == 'stacked':
from datasets.stacked_sequence_folders import SequenceFolder
elif args.dataset_format == 'sequential':
from datasets.sequence_folders import SequenceFolder, StereoSequenceFolder
save_path = save_path_formatter(args, parser)
args.save_path = 'checkpoints'/save_path
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
if args.evaluate:
args.epochs = 0
training_writer = SummaryWriter(args.save_path)
output_writers = []
if args.log_output:
for i in range(3):
output_writers.append(SummaryWriter(args.save_path/'valid'/str(i)))
# Data loading code
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(),
normalize
])
valid_transform = custom_transforms.Compose([custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = StereoSequenceFolder(
args.data,
transform=train_transform,
seed=args.seed,
train=True,
sequence_length=args.sequence_length
)
# if no Groundtruth is avalaible, Validation set is the same type as training set to measure photometric loss from warping
if args.with_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(
args.data,
transform=valid_transform
)
else:
val_set = StereoSequenceFolder(
args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
)
print('{} samples found in {} train scenes'.format(len(train_set), len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set), len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
val_set, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
# 没有epoch_size的时候（=0），每个epoch训练train_set中所有的samples
# 有epoch_size的时候，每个epoch只训练一部分train_set
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
# 初始化网络结构
print("=> creating model")
# disp_net = models.DispNetS().to(device)
disp_net = models.DispResNet(3).to(device)
output_exp = args.mask_loss_weight > 0
if not output_exp:
print("=> no mask loss, PoseExpnet will only output pose")
# 如果有mask loss，PoseExpNet 要输出mask和pose estimation，因为两个输出共享encoder网络
# pose_exp_net = PoseExpNet(nb_ref_imgs=args.sequence_length - 1, output_exp=args.mask_loss_weight > 0).to(device)
pose_exp_net = models.PoseExpNet(nb_ref_imgs=args.sequence_length - 1, output_exp=args.mask_loss_weight > 0).to(device)
if args.pretrained_exp_pose:
print("=> using pre-trained weights for explainabilty and pose net")
weights = torch.load(args.pretrained_exp_pose)
pose_exp_net.load_state_dict(weights['state_dict'], strict=False)
else:
pose_exp_net.init_weights()
if args.pretrained_disp:
print("=> using pre-trained weights for Dispnet")
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'])
else:
disp_net.init_weights()
cudnn.benchmark = True
# 并行化
disp_net = torch.nn.DataParallel(disp_net)
pose_exp_net = torch.nn.DataParallel(pose_exp_net)
# 训练方式：Adam
print('=> setting adam solver')
# 两个网络一起
optim_params = [
{'params': disp_net.parameters(), 'lr': args.lr},
{'params': pose_exp_net.parameters(), 'lr': args.lr}
]
optimizer = torch.optim.Adam(optim_params,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
with open(args.save_path/args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path/args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'photo_loss', 'explainability_loss', 'smooth_loss'])
# 对pretrained模型先做评估
if args.pretrained_disp or args.evaluate:
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net, 0, output_writers)
else:
errors, error_names = validate_without_gt(args, val_loader, disp_net, pose_exp_net, 0, output_writers)
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error, 0)
error_string = ', '.join('{} : {:.3f}'.format(name, error) for name, error in zip(error_names[2:9], errors[2:9]))
# 正式训练
for epoch in range(args.epochs):
# train for one epoch 训练一个周期
print('\n')
train_loss = train(args, train_loader, disp_net, pose_exp_net, optimizer, args.epoch_size, training_writer, epoch)
# evaluate on validation set
print('\n')
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net, epoch, output_writers)
else:
errors, error_names = validate_without_gt(args, val_loader, disp_net, pose_exp_net, epoch, output_writers)
error_string = ', '.join('{} : {:.3f}'.format(name, error) for name, error in zip(error_names, errors))
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error, epoch)
# Up to you to chose the most relevant error to measure your model's performance, careful some measures are to maximize (such as a1,a2,a3)
# 验证输出四个loss：总体final loss，warping loss以及mask正则化loss
# 可自选以哪一种loss作为best model的标准
decisive_error = errors[0]
if best_error < 0:
best_error = decisive_error
# remember lowest error and save checkpoint
# 保存validation最佳model
is_best = decisive_error < best_error
best_error = min(best_error, decisive_error)
save_checkpoint(
args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_exp_net.module.state_dict()
},
is_best)
with open(args.save_path/args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])
def main():
global args, best_error, n_iter
if args.dataset_format == 'stacked':
from datasets.stacked_sequence_folders import SequenceFolder
elif args.dataset_format == 'sequential':
from datasets.sequence_folders import SequenceFolder
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
args.save_path = Path('checkpoints') / args.name / timestamp
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
if args.alternating:
args.alternating_flags = np.array([False, False, True])
training_writer = SummaryWriter(args.save_path)
output_writers = []
if args.log_output:
for i in range(3):
output_writers.append(
SummaryWriter(args.save_path / 'valid' / str(i)))
# Data loading code
flow_loader_h, flow_loader_w = 256, 832
if args.data_normalization == 'global':
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
elif args.data_normalization == 'local':
normalize = custom_transforms.NormalizeLocally()
if args.fix_flownet:
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
else:
train_transform = custom_transforms.Compose([
custom_transforms.RandomRotate(),
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
valid_flow_transform = custom_transforms.Compose([
custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w),
custom_transforms.ArrayToTensor(), normalize
])
print("=> fetching scenes in '{}'".format(args.data))
train_set = SequenceFolder(args.data,
transform=train_transform,
seed=args.seed,
train=True,
sequence_length=args.sequence_length)
# if no Groundtruth is avalaible, Validation set is the same type as training set to measure photometric loss from warping
if args.with_depth_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(args.data.replace('cityscapes', 'kitti'),
transform=valid_transform)
else:
val_set = SequenceFolder( #只有图
args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
)
if args.with_flow_gt:
from datasets.validation_flow import ValidationFlow
val_flow_set = ValidationFlow(root=args.kitti_dir,
sequence_length=args.sequence_length,
transform=valid_flow_transform)
val_flow_loader = torch.utils.data.DataLoader(
val_flow_set,
batch_size=1,
# batch size is 1 since images in kitti have different sizes
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True)
if args.DEBUG:
train_set.__len__ = 32
train_set.samples = train_set.samples[:32]
print('{} samples found in {} train scenes'.format(len(train_set),
len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
#1 create model
print("=> creating model")
#1.1 disp_net
disp_net = getattr(models, args.dispnet)().cuda()
output_exp = True #args.mask_loss_weight > 0
if not output_exp:
print("=> no mask loss, PoseExpnet will only output pose")
#1.2 pose_net
pose_net = getattr(models, args.posenet)(nb_ref_imgs=args.sequence_length -
1).cuda()
#1.3.flow_net
if args.flownet == 'SpyNet':
flow_net = getattr(models,
args.flownet)(nlevels=args.nlevels,
pre_normalization=normalize).cuda()
elif args.flownet == 'FlowNetC6': #flonwtc6
flow_net = getattr(models, args.flownet)(nlevels=args.nlevels).cuda()
elif args.flownet == 'FlowNetS':
flow_net = getattr(models, args.flownet)(nlevels=args.nlevels).cuda()
elif args.flownet == 'Back2Future':
flow_net = getattr(models, args.flownet)(nlevels=args.nlevels).cuda()
# 1.4 mask_net
mask_net = getattr(models,
args.masknet)(nb_ref_imgs=args.sequence_length - 1,
output_exp=True).cuda()
#2 载入参数
#2.1 pose
if args.pretrained_pose:
print("=> using pre-trained weights for explainabilty and pose net")
weights = torch.load(args.pretrained_pose)
pose_net.load_state_dict(weights['state_dict'])
else:
pose_net.init_weights()
if args.pretrained_mask:
print("=> using pre-trained weights for explainabilty and pose net")
weights = torch.load(args.pretrained_mask)
mask_net.load_state_dict(weights['state_dict'])
else:
mask_net.init_weights()
# import ipdb; ipdb.set_trace()
if args.pretrained_disp:
print("=> using pre-trained weights from {}".format(
args.pretrained_disp))
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'])
else:
disp_net.init_weights()
if args.pretrained_flow:
print("=> using pre-trained weights for FlowNet")
weights = torch.load(args.pretrained_flow)
flow_net.load_state_dict(weights['state_dict'])
else:
flow_net.init_weights()
if args.resume:
print("=> resuming from checkpoint")
dispnet_weights = torch.load(args.save_path /
'dispnet_checkpoint.pth.tar')
posenet_weights = torch.load(args.save_path /
'posenet_checkpoint.pth.tar')
masknet_weights = torch.load(args.save_path /
'masknet_checkpoint.pth.tar')
flownet_weights = torch.load(args.save_path /
'flownet_checkpoint.pth.tar')
disp_net.load_state_dict(dispnet_weights['state_dict'])
pose_net.load_state_dict(posenet_weights['state_dict'])
flow_net.load_state_dict(flownet_weights['state_dict'])
mask_net.load_state_dict(masknet_weights['state_dict'])
# import ipdb; ipdb.set_trace()
cudnn.benchmark = True
disp_net = torch.nn.DataParallel(disp_net)
pose_net = torch.nn.DataParallel(pose_net)
mask_net = torch.nn.DataParallel(mask_net)
flow_net = torch.nn.DataParallel(flow_net)
print('=> setting adam solver')
parameters = chain(disp_net.parameters(), pose_net.parameters(),
mask_net.parameters(), flow_net.parameters())
optimizer = torch.optim.Adam(parameters,
args.lr,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
if args.resume and (args.save_path /
'optimizer_checkpoint.pth.tar').exists():
print("=> loading optimizer from checkpoint")
optimizer_weights = torch.load(args.save_path /
'optimizer_checkpoint.pth.tar')
optimizer.load_state_dict(optimizer_weights['state_dict'])
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([
'train_loss', 'photo_cam_loss', 'photo_flow_loss',
'explainability_loss', 'smooth_loss'
])
if args.log_terminal:
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
else:
logger = None
#3. main cycle
for epoch in range(args.epochs):
#3.1 四个子网络，训练哪几个
if args.fix_flownet:
for fparams in flow_net.parameters():
fparams.requires_grad = False
if args.fix_masknet:
for fparams in mask_net.parameters():
fparams.requires_grad = False
if args.fix_posenet:
for fparams in pose_net.parameters():
fparams.requires_grad = False
if args.fix_dispnet:
for fparams in disp_net.parameters():
fparams.requires_grad = False
if args.log_terminal:
logger.epoch_bar.update(epoch)
logger.reset_train_bar()
#validation data
flow_error_names = ['no']
flow_errors = [0]
errors = [0]
error_names = ['no error names depth']
print('\nepoch [{}/{}]\n'.format(epoch + 1, args.epochs))
#3.2 train for one epoch---------
train_loss = 0
train_loss = train(train_loader, disp_net, pose_net, mask_net,
flow_net, optimizer, logger, training_writer)
#3.3 evaluate on validation set-----
#if args.without_gt:
# validate_without_gt(val_loader,disp_net,pose_net,mask_net, epoch, logger, output_writers)
if args.with_flow_gt:
flow_errors, flow_error_names = validate_flow_with_gt(
val_flow_loader, disp_net, pose_net, mask_net, flow_net, epoch,
logger, output_writers)
for error, name in zip(flow_errors, flow_error_names):
training_writer.add_scalar(name, error, epoch)
if args.with_depth_gt:
depth_errors, depth_error_names = validate_depth_with_gt(
val_loader, disp_net, epoch, logger, output_writers)
error_string = ', '.join(
'{} : {:.3f}'.format(name, error)
for name, error in zip(depth_error_names, depth_errors))
if args.log_terminal:
logger.valid_writer.write(' * Avg {}'.format(error_string))
else:
print('Epoch {} completed'.format(epoch))
for error, name in zip(depth_errors, depth_error_names):
training_writer.add_scalar(name, error, epoch)
if args.without_gt:
val_loss = 1
#----------------------
#3.4 Up to you to chose the most relevant error to measure your model's performance, careful some measures are to maximize (such as a1,a2,a3)
if not args.fix_posenet:
decisive_error = 0 # flow_errors[-2] # epe_rigid_with_gt_mask
elif not args.fix_dispnet:
decisive_error = 0 # errors[0] #depth abs_diff
elif not args.fix_flownet:
decisive_error = 0 # flow_errors[-1] #epe_non_rigid_with_gt_mask
elif not args.fix_masknet:
decisive_error = 0 #flow_errors[3] # percent outliers
#3.5 log
if args.log_terminal:
logger.train_writer.write(
' * Avg Loss : {:.3f}'.format(train_loss))
logger.reset_valid_bar()
#3.6 save model and remember lowest error and save checkpoint
if best_error < 0:
best_error = decisive_error
is_best = decisive_error <= best_error
best_error = min(best_error, decisive_error)
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': mask_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': flow_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': optimizer.state_dict()
}, is_best)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])
if args.log_terminal:
logger.epoch_bar.finish()
def main():
global best_error, n_iter, device
args = parser.parse_args()
if args.dataset_format == 'stacked':
from datasets.stacked_sequence_folders import SequenceFolder
elif args.dataset_format == 'sequential':
from datasets.sequence_folders import SequenceFolder
save_path = save_path_formatter(args, parser)
args.save_path = 'checkpoints' / save_path
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p() #如果没有，则建立，有则啥都不干 in Path.py小工具
torch.manual_seed(args.seed)
if args.evaluate:
args.epochs = 0
#tensorboard SummaryWriter
training_writer = SummaryWriter(args.save_path) #for tensorboard
output_writers = [] #list
if args.log_output:
for i in range(3):
output_writers.append(
SummaryWriter(args.save_path / 'valid' / str(i)))
# Data loading code
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
'''transform'''
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = SequenceFolder(
args.data, #processed_data_train_sets
transform=train_transform, #把几种变换函数输入进去
seed=args.seed,
train=True,
sequence_length=args.sequence_length)
# if no Groundtruth is avalaible, Validation set is
# the same type as training set to measure photometric loss from warping
if args.with_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(args.data, transform=valid_transform)
else:
val_set = SequenceFolder(
args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
)
print('{} samples found in {} train scenes'.format(
len(train_set), len(train_set.scenes))) #训练集都是序列,不用左右
print('{} samples found in {} valid scenes'.format(
len(val_set), len(val_set.scenes))) #测试集也是序列,不需要左右
train_loader = torch.utils.data.DataLoader( #data(list): [tensor(B,3,H,W),list(B),(B,H,W),(b,h,w)]
dataset=train_set, #sequenceFolder
batch_size=args.batch_size,
shuffle=True, #打乱
num_workers=args.workers, #多线程读取数据
pin_memory=True)
val_loader = torch.utils.data.DataLoader(
dataset=val_set,
batch_size=args.batch_size,
shuffle=False, #不打乱
num_workers=args.workers,
pin_memory=True)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
print("=> creating model")
#disp
disp_net = models.DispNetS().to(device)
output_exp = args.mask_loss_weight > 0
if not output_exp:
print("=> no mask loss, PoseExpnet will only output pose")
#pose
pose_exp_net = models.PoseExpNet(
nb_ref_imgs=args.sequence_length - 1,
output_exp=args.mask_loss_weight > 0).to(device)
#init posenet
if args.pretrained_exp_pose:
print("=> using pre-trained weights for explainabilty and pose net")
weights = torch.load(args.pretrained_exp_pose)
pose_exp_net.load_state_dict(weights['state_dict'], strict=False)
else:
pose_exp_net.init_weights()
#init dispNet
if args.pretrained_disp:
print("=> using pre-trained weights for Dispnet")
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'])
else:
disp_net.init_weights()
cudnn.benchmark = True
disp_net = torch.nn.DataParallel(disp_net)
pose_exp_net = torch.nn.DataParallel(pose_exp_net)
print('=> setting adam solver')
#可以看到两个一起训练
optim_params = [{
'params': disp_net.parameters(),
'lr': args.lr
}, {
'params': pose_exp_net.parameters(),
'lr': args.lr
}]
optimizer = torch.optim.Adam(optim_params,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
#训练结果写入csv
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(
['train_loss', 'photo_loss', 'explainability_loss', 'smooth_loss'])
n_epochs = args.epochs
train_size = min(len(train_loader), args.epoch_size)
valid_size = len(val_loader)
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
if args.pretrained_disp or args.evaluate:
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
0, logger, output_writers)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_exp_net,
0, logger,
output_writers)
for error, name in zip(
errors, error_names
): #validation时，对['Total loss', 'Photo loss', 'Exp loss']三个 epoch-record 指标添加记录值
training_writer.add_scalar(name, error, 0)
error_string = ', '.join(
'{} : {:.3f}'.format(name, error)
for name, error in zip(error_names[2:9], errors[2:9]))
logger.valid_writer.write(' * Avg {}'.format(error_string))
#main cycle
for epoch in range(args.epochs):
logger.epoch_bar.update(epoch)
logger.reset_train_bar()
#1. train for one epoch
train_loss = train(args, train_loader, disp_net, pose_exp_net,
optimizer, args.epoch_size, logger, training_writer)
#其他参数都好解释， logger: SelfDefined class,
logger.train_writer.write(' * Avg Loss : {:.3f}'.format(train_loss))
logger.reset_valid_bar()
# 2. validate on validation set
if args.with_gt: #<class 'list'>: ['Total loss', 'Photo loss', 'Exp loss']
errors, error_names = validate_with_gt(args, val_loader, disp_net,
epoch, logger,
output_writers)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_exp_net,
epoch, logger,
output_writers)
error_string = ', '.join('{} : {:.3f}'.format(name, error)
for name, error in zip(error_names, errors))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error,
epoch) #损失函数中记录epoch-record指标
# Up to you to chose the most relevant error to measure
# your model's performance, careful some measures are to maximize (such as a1,a2,a3)
# 3. remember lowest error and save checkpoint
decisive_error = errors[1]
if best_error < 0:
best_error = decisive_error
is_best = decisive_error < best_error
best_error = min(best_error, decisive_error)
#模型保存
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_exp_net.module.state_dict()
}, is_best)
with open(args.save_path / args.log_summary,
'a') as csvfile: #每个epoch留下结果
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss,
decisive_error]) #第二个就是validataion 中的epoch-record
# loss<class 'list'>: ['Total loss', 'Photo loss', 'Exp loss']
logger.epoch_bar.finish()
