前言
在系列七中,我們提到了train.py中實際上只有兩行訓練相關的代碼,第一行是Trainer構造函數的調用,主要是初始化和數據集的構建,系列七主要是對這個過程進行了梳理。第二行是Trainer成員函數train的執行,這個是訓練真正執行部分,本文着重來對它進行分析。
訓練數據集按batch來加載
def train(self):
"""Run the entire training pipeline
"""
self.epoch = 0
self.step = 0
self.start_time = time.time()
for self.epoch in range(self.opt.num_epochs):
self.run_epoch()
if (self.epoch + 1) % self.opt.save_frequency == 0:
self.save_model()
上面是成員函數train的實現,主要部分在run_epoch()裏面。
def run_epoch(self):
"""Run a single epoch of training and validation
"""
#self.model_lr_scheduler.step()
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)
。。。 。。。
這個函數其實也很簡單,首先通過set_train()來將resnet encoder和depth decoder模型設置成訓練狀態,然後通過enumrate(self.train_loader)來返回一個batch大小的inputs數據。
上文說過,每當枚舉train_loader時,就會調用一次mono_dataset.py中的__getitem__()。這個函數較複雜些,但很重要。方便起見,我在代碼裏面添加了中文解釋信息。
def __getitem__(self, index):
inputs = {}
//隨機做訓練數據顏色增強預處理
do_color_aug = self.is_train and random.random() > 0.5
//隨機做訓練數據水平左右flip預處理
do_flip = self.is_train and random.random() > 0.5
//index是train_txt中的第index行。
line = self.filenames[index].split()
//train_files.txt中一行數據的第一部分,即圖片所在目錄。
folder = line[0]
//每一行一般都爲3個部分,第二個部分是圖片的frame_index
if len(line) == 3:
frame_index = int(line[1])
else:
frame_index = 0
//side爲l或r,表明該圖片是左或右攝像頭所拍。
if len(line) == 3:
side = line[2]
else:
side = None
//在stereo訓練時, frame_idxs爲["0","s"]
//通過這個for循環,inputs[("color", "0", -1)]和inputs[("color", "s", -1)]
//分別獲得了frame_index和它對應的另外一個攝像頭拍的圖片數據。
for i in self.frame_idxs:
if i == "s":
other_side = {"r": "l", "l": "r"}[side]
inputs[("color", i, -1)] = self.get_color(folder, frame_index, other_side, do_flip)
else:
inputs[("color", i, -1)] = self.get_color(folder, frame_index + i, side, do_flip)
# adjusting intrinsics to match each scale in the pyramid
//因爲模型有4個尺度,所以對應4個相機內參
for scale in range(self.num_scales):
K = self.K.copy()
K[0, :] *= self.width // (2 ** scale)
K[1, :] *= self.height // (2 ** scale)
inv_K = np.linalg.pinv(K)
inputs[("K", scale)] = torch.from_numpy(K)
inputs[("inv_K", scale)] = torch.from_numpy(inv_K)
//顏色增強參數設定
if do_color_aug:
color_aug = transforms.ColorJitter.get_params(
self.brightness, self.contrast, self.saturation, self.hue)
else:
color_aug = (lambda x: x)
//訓練前數據預處理以及對輸入數據做多尺度resize。
self.preprocess(inputs, color_aug)
//經過preprocess,產生了inputs[("color","0", 0/1/23)]和inputs[("color_aug","0",
// 0/1/23)]。所以可以將原始的inputs[("color", i, -1)]和[("color_aug", i, -1)]釋放
for i in self.frame_idxs:
del inputs[("color", i, -1)]
del inputs[("color_aug", i, -1)]
//load_depth爲False,因爲不需要GT label數據
if self.load_depth:
depth_gt = self.get_depth(folder, frame_index, side, do_flip)
inputs["depth_gt"] = np.expand_dims(depth_gt, 0)
inputs["depth_gt"] = torch.from_numpy(inputs["depth_gt"].astype(np.float32))
//在stereo訓練時,還需要構造雙目姿態的平移矩陣參數inputs["stereo_T"]
if "s" in self.frame_idxs:
stereo_T = np.eye(4, dtype=np.float32)
baseline_sign = -1 if do_flip else 1
side_sign = -1 if side == "l" else 1
stereo_T[0, 3] = side_sign * baseline_sign * 0.1
inputs["stereo_T"] = torch.from_numpy(stereo_T)
return inputs
開始處理
通過上面的枚舉train_loader操作就可以得到各個尺度的inputs數據,然後作爲參數輸入到self.process_batch(inputs)。process_batch的返回值爲ouputs和loss。這個函數執行完後,整個train就只剩下根據loss值backward來更新梯度,並根據優化器和lr來更新權值。
def process_batch(self, inputs):
"""Pass a minibatch through the network and generate images and losses
"""
for key, ipt in inputs.items():
inputs[key] = ipt.to(self.device)
if self.opt.pose_model_type == "shared":
# If we are using a shared encoder for both depth and pose (as advocated
# in monodepthv1), then all images are fed separately through the depth encoder.
all_color_aug = torch.cat([inputs[("color_aug", i, 0)] for i in self.opt.frame_ids])
all_features = self.models["encoder"](all_color_aug)
all_features = [torch.split(f, self.opt.batch_size) for f in all_features]
features = {}
for i, k in enumerate(self.opt.frame_ids):
features[k] = [f[i] for f in all_features]
outputs = self.models["depth"](features[0])
else:
# Otherwise, we only feed the image with frame_id 0 through the depth encoder
features = self.models["encoder"](inputs["color_aug", 0, 0])
outputs = self.models["depth"](features)
if self.opt.predictive_mask:
outputs["predictive_mask"] = self.models["predictive_mask"](features)
if self.use_pose_net:
outputs.update(self.predict_poses(inputs, features))
self.generate_images_pred(inputs, outputs)
losses = self.compute_losses(inputs, outputs)
return outputs, losses
在上面的函數中,outputs是depth decoder求出來的,具體代碼爲:
features = self.models["encoder"](inputs["color_aug", 0, 0])和outputs = self.models["depth"](features)。
有了ouputs就可以來算loss,這個主要通過self.generate_images_pred(inputs, outputs)和losses = self.compute_losses(inputs, outputs)來實現。 細節將在下一篇文章來分析。