雙目測距系列(八)monodepth2訓練代碼分析上

前言

在系列七中,我們提到了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)來實現。 細節將在下一篇文章來分析。


 

 


 

 

 

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章