参考代码:manydepth
论文名称:The Temporal Opportunist: Self-Supervised Multi-Frame Monocular Depth
manydepth项目中train部分
if __name__ == "__main__":
trainer = Trainer(opts)#实例化
trainer.train()
初始化函数,初始化函数的主要作用是引入参数,获得路径
def __init__(self, options):
self.opt = options
self.log_path = os.path.join(self.opt.log_dir, self.opt.model_name)
#合成的路径,把逗号隔开的连在一起
assert self.opt.height % 32 == 0, "'height' must be a multiple of 32"
assert self.opt.width % 32 == 0, "'width' must be a multiple of 32"
#判断是否是32的整数倍,改变H,W的尺度
self.models = {}
self.parameters_to_train = []
self.device = torch.device("cpu" if self.opt.no_cuda else "cuda")
self.num_scales = len(self.opt.scales)#尺度数量
self.num_input_frames = len(self.opt.frame_ids)#输入帧的数量
self.num_pose_frames = 2#姿态帧的数量
assert self.opt.frame_ids[0] == 0,
assert len(self.opt.frame_ids) > 1,
#frame_ids必须从0开始,有1个以上的指定帧
self.train_teacher_and_pose = not self.opt.freeze_teacher_and_pose
if self.train_teacher_and_pose:#相机估计位姿网络
print('using adaptive depth binning!')#打印自适应深度合并
self.min_depth_tracker = 0.1#深度范围0.1-10
self.max_depth_tracker = 10.0
else:
print('fixing pose network and monocular network!')
#检查所需数据加载的帧
frames_to_load = self.opt.frame_ids.copy()
self.matching_ids = [0]
if self.opt.use_future_frame:#使用未来帧作为监督信号
self.matching_ids.append(1)
for idx in range(-1, -1 - self.opt.num_matching_frames, -1):#遍历索引
self.matching_ids.append(idx)
if idx not in frames_to_load:
frames_to_load.append(idx)
print('Loading frames: {}'.format(frames_to_load))#输出加载的帧
建立model ,引入network文件夹中ResnetEncoderMatching、DepthDecoder、ResnetEncoder、PoseDecoder几个类。
Resnet编码器适应于在第2个区块后包括cost-volume。adaptive_bins=True:在每次向前传递时,用于匹配的深度分量将被重新计算,这对于单目视频的训练是必需的,因为有一个未知的尺度。PoseDecoder:输入帧进入posecnn,由7层conv组成,PoseDecoder由3层conv组成。
self.models["encoder"] = networks.ResnetEncoderMatching(
self.opt.num_layers, self.opt.weights_init == "pretrained",
input_height=self.opt.height, input_width=self.opt.width,
adaptive_bins=True, min_depth_bin=0.1, max_depth_bin=20.0,
depth_binning=self.opt.depth_binning, num_depth_bins=self.opt.num_depth_bins)
self.models["encoder"].to(self.device)
self.models["depth"] = networks.DepthDecoder(
self.models["encoder"].num_ch_enc, self.opt.scales)
self.models["depth"].to(self.device)
self.parameters_to_train += list(self.models["encoder"].parameters())
self.parameters_to_train += list(self.models["depth"].parameters())
self.models["mono_encoder"] =
networks.ResnetEncoder(18, self.opt.weights_init == "pretrained")
self.models["mono_encoder"].to(self.device)
self.models["mono_depth"] =
networks.DepthDecoder(self.models["mono_encoder"].num_ch_enc, self.opt.scales)
self.models["mono_depth"].to(self.device)
if self.train_teacher_and_pose:
self.parameters_to_train += list(self.models["mono_encoder"].parameters())
self.parameters_to_train += list(self.models["mono_depth"].parameters())
self.models["pose_encoder"] =
networks.ResnetEncoder(18, self.opt.weights_init == "pretrained",
num_input_images=self.num_pose_frames)#resnet18
self.models["pose_encoder"].to(self.device)
self.models["pose"] =
networks.PoseDecoder(self.models["pose_encoder"].num_ch_enc,
num_input_features=1,
num_frames_to_predict_for=2)
self.models["pose"].to(self.device)
if self.train_teacher_and_pose:
self.parameters_to_train += list(self.models["pose_encoder"].parameters())
self.parameters_to_train += list(self.models["pose"].parameters())
self.model_optimizer = optim.Adam(self.parameters_to_train, self.opt.learning_rate)
self.model_lr_scheduler = optim.lr_scheduler.StepLR(
self.model_optimizer, self.opt.scheduler_step_size, 0.1)
if self.opt.load_weights_folder is not None:
self.load_model()
if self.opt.mono_weights_folder is not None:
self.load_mono_model()
print("Training model named:n ", self.opt.model_name)
print("Models and tensorboard events files are saved to:n ", self.opt.log_dir)
print("Training is using:n ", self.device)
引入数据集
# DATA
datasets_dict = {"kitti": datasets.KITTIRAWDataset,
"cityscapes_preprocessed": datasets.CityscapesPreprocessedDataset,
"kitti_odom": datasets.KITTIOdomDataset}
self.dataset = datasets_dict[self.opt.dataset]
fpath = os.path.join("../splits", self.opt.split, "{}_files.txt"#切分训练集和验证集
train_filenames = readlines(fpath.format("train"))#设置训练集文件夹名
val_filenames = readlines(fpath.format("val"))#设置验证集文件夹名
img_ext = '.png' if self.opt.png else '.jpg'
num_train_samples = len(train_filenames)#训练样本数量
self.num_total_steps = num_train_samples // self.opt.batch_size * self.opt.num_epochs
train_dataset = self.dataset(
self.opt.data_path, train_filenames, self.opt.height, self.opt.width,
frames_to_load, 4, is_train=True, img_ext=img_ext)
self.train_loader = DataLoader(
train_dataset, self.opt.batch_size, True,
num_workers=self.opt.num_workers, pin_memory=True, drop_last=True,
worker_init_fn=seed_worker)
val_dataset = self.dataset(
self.opt.data_path, val_filenames, self.opt.height, self.opt.width,
frames_to_load, 4, is_train=False, img_ext=img_ext)
self.val_loader = DataLoader(
val_dataset, self.opt.batch_size, True,
num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)
self.val_iter = iter(self.val_loader)
self.writers = {}
for mode in ["train", "val"]:
self.writers[mode] = SummaryWriter(os.path.join(self.log_path, mode))
if not self.opt.no_ssim:
self.ssim = SSIM()
self.ssim.to(self.device)
self.backproject_depth = {}
self.project_3d = {}
for scale in self.opt.scales:
h = self.opt.height // (2 ** scale)
w = self.opt.width // (2 ** scale)
self.backproject_depth[scale] = BackprojectDepth(self.opt.batch_size, h, w)
self.backproject_depth[scale].to(self.device)
self.project_3d[scale] = Project3D(self.opt.batch_size, h, w)
self.project_3d[scale].to(self.device)
self.depth_metric_names = [
"de/abs_rel", "de/sq_rel", "de/rms", "de/log_rms", "da/a1", "da/a2", "da/a3"]
print("Using split:n ", self.opt.split)
print("There are {:d} training items and {:d} validation itemsn".format(
len(train_dataset), len(val_dataset)))
self.save_opts()
将所有的model转换训练模式
def set_train(self):
for k, m in self.models.items():
if self.train_teacher_and_pose:
m.train()
else:
# if teacher + pose is frozen, then only use training batch norm stats for
# multi components
if k in ['depth', 'encoder']:
m.train()
将所有模型转换为测试/评估模式
def set_eval(self):
for m in self.models.values():
m.eval()
运行整个训练通道
def train(self):
self.epoch = 0
self.step = 0
self.start_time = time.time()#开始计时
for self.epoch in range(self.opt.num_epochs):#循环epoch
if self.epoch == self.opt.freeze_teacher_epoch:
self.freeze_teacher()
self.run_epoch()
if (self.epoch + 1) % self.opt.save_frequency == 0:
self.save_model()
在自监督学习时,freeze预训练模型,不再更新其参数。
def freeze_teacher(self):
if self.train_teacher_and_pose:
self.train_teacher_and_pose = False
print('freezing teacher and pose networks!')
# here we reinitialise our optimizer to ensure there are no updates to the
# teacher and pose networks
self.parameters_to_train = []
self.parameters_to_train += list(self.models["encoder"].parameters())
self.parameters_to_train += list(self.models["depth"].parameters())
self.model_optimizer = optim.Adam(self.parameters_to_train, self.opt.learning_rate)
self.model_lr_scheduler = optim.lr_scheduler.StepLR(
self.model_optimizer, self.opt.scheduler_step_size, 0.1)
# set eval so that teacher + pose batch norm is running average
self.set_eval()
# set train so that multi batch norm is in train mode
self.set_train()
运行单个的训练和验证epoch
def run_epoch(self):
"""Run a single epoch of training and validation
"""
print("Training")
self.set_train()
for batch_idx, inputs in enumerate(self.train_loader):
before_op_time = time.time()
outputs, losses = self.process_batch(inputs, is_train=True)
self.model_optimizer.zero_grad()
losses["loss"].backward()
self.model_optimizer.step()
duration = time.time() - before_op_time
# log less frequently after the first 2000 steps to save time & disk space
early_phase = batch_idx % self.opt.log_frequency == 0 and self.step < 2000
late_phase = self.step % 2000 == 0
if early_phase or late_phase:
self.log_time(batch_idx, duration, losses["loss"].cpu().data)
if "depth_gt" in inputs:
self.compute_depth_losses(inputs, outputs, losses)
self.log("train", inputs, outputs, losses)
self.val()
if self.opt.save_intermediate_models and late_phase:
self.save_model(save_step=True)
if self.step == self.opt.freeze_teacher_step:
self.freeze_teacher()
self.step += 1
self.model_lr_scheduler.step()



