trainer_icdepth.py

# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import numpy as np
import time

import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
import pdb
import json

from utils import *
from kitti_utils import *
from layers import *
#import save
import datasets
import networks
from IPython import embed


class Trainer:
    def __init__(self, options):
        self.h = 0
        self.opt = options
        self.log_path = os.path.join(self.opt.log_dir, self.opt.model_name)

        # checking height and width are multiples of 32
        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"

        self.models = {}
        self.model = {}
        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 if self.opt.pose_model_input == "pairs" else self.num_input_frames

        assert self.opt.frame_ids[0] == 0, "frame_ids must start with 0"

        self.use_pose_net = not (self.opt.use_stereo and self.opt.frame_ids == [0])

        if self.opt.use_stereo:
            self.opt.frame_ids.append("s")

        self.models["encoder"] = networks.mpvit_small()
        self.models["encoder"].num_ch_enc = [64, 128, 216, 288, 288]
        self.models["encoder"] = nn.DataParallel(self.models["encoder"])
        self.models["encoder"].to(self.device)
        # self.parameters_to_train += list(self.models["encoder"].parameters())

        self.models["depth"] = networks.DepthDecoder(
            self.models["encoder"].module.num_ch_enc, self.opt.scales)
        self.models["depth"] = nn.DataParallel(self.models["depth"])
        self.models["depth"].to(self.device)
        self.parameters_to_train += list(self.models["depth"].parameters())

        if self.use_pose_net:
            if self.opt.pose_model_type == "separate_resnet":
                self.model["pose_encoder"] = networks.ResnetEncoder(
                    self.opt.num_layers,
                    self.opt.weights_init == "pretrained",
                    num_input_images=self.num_pose_frames)
                self.model["pose_encoder"] = nn.DataParallel(self.model["pose_encoder"])
                self.model["pose_encoder"].to(self.device)
                self.parameters_to_train += list(self.model["pose_encoder"].parameters())

                self.models["pose"] = networks.PoseDecoder(
                    self.model["pose_encoder"].module.num_ch_enc,
                    num_input_features=1,
                    num_frames_to_predict_for=2)

            elif self.opt.pose_model_type == "shared":
                self.models["pose"] = networks.PoseDecoder(
                    self.models["encoder"].module.num_ch_enc, self.num_pose_frames)

            elif self.opt.pose_model_type == "posecnn":
                self.models["pose"] = networks.PoseCNN(
                    self.num_input_frames if self.opt.pose_model_input == "all" else 2)
            self.models["pose"] = nn.DataParallel(self.models["pose"])
            self.models["pose"].to(self.device)
            self.parameters_to_train += list(self.models["pose"].parameters())

        if self.opt.predictive_mask:
            assert self.opt.disable_automasking, \
                "When using predictive_mask, please disable automasking with --disable_automasking"

            # Our implementation of the predictive masking baseline has the the same architecture
            # as our depth decoder. We predict a separate mask for each source frame.
            self.models["predictive_mask"] = networks.DepthDecoder(
                self.models["encoder"].num_ch_enc, self.opt.scales,
                num_output_channels=(len(self.opt.frame_ids) - 1))
            self.models["predictive_mask"].to(self.device)
            self.parameters_to_train += list(self.models["predictive_mask"].parameters())


        self.params = [ {
            "params":self.parameters_to_train,
            "lr": self.opt.learning_rate
            #"weight_decay": 0.01
            },
            {
            "params": list(self.models["encoder"].parameters()),
           "lr": self.opt.learning_rate/2
            #"weight_decay": 0.01
            } ]
        self.model_optimizer = optim.AdamW(self.params)
        self.model_lr_scheduler = optim.lr_scheduler.ExponentialLR(self.model_optimizer,0.9)


        if self.opt.load_weights_folder is not None:
            self.load_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)
        if not os.path.exists(self.log_path):
            os.makedirs(self.log_path)
        save_log = os.path.join(self.log_path,'results1.txt')
        results_edit = open(save_log, mode='a')
        results_edit.write("Training model named:\n  "+self.opt.model_name)
        results_edit.write("Models and tensorboard events files are saved to:\n  "+self.opt.log_dir)
        results_edit.write("Training is using:\n  "+"cuda")
        results_edit.close()

        # data
        datasets_dict = {"kitti": datasets.KITTIRAWDataset,
                         "kitti_odom": datasets.KITTIOdomDataset}
        self.dataset = datasets_dict[self.opt.dataset]

        fpath = os.path.join(os.path.dirname(__file__), "splits", self.opt.split, "{}_files.txt")


        # train_filenames = readlines(fpath.format("train"))
        train_filenames = readlines(fpath.format("train1"))

        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,
            self.opt.frame_ids, 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)
        val_dataset = self.dataset(
            self.opt.data_path, val_filenames, self.opt.height, self.opt.width,
            self.opt.frame_ids, 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 items\n".format(
            len(train_dataset), len(val_dataset)))
        save_log = os.path.join(self.log_path,'results1.txt')
        results_edit = open(save_log, mode='a')
        results_edit.write("Using split:\n  "+self.opt.split)
        results_edit.write("There are {:d} training items and {:d} validation items\n".format(
            len(train_dataset), len(val_dataset)))
        results_edit.close()

        self.save_opts()

    def set_train(self):
        """Convert all models to training mode
        """
        for m in self.models.values():
            m.train()

    def set_eval(self):
        """Convert all models to testing/evaluation mode
        """
        for m in self.models.values():
            m.eval()

    def train(self):
        """Run the entire training pipeline
        """
        self.epoch = 0
        self.step = 0
        self.dep_wight1=0
        self.image_wight1 = 0
        self.start_time = time.time()
        #pdb.set_trace()
        for self.epoch in range(self.opt.num_epochs):
            # if self.epoch == 1:
            #     for params in self.model_optimizer.param_groups:
            #         params['lr'] = 1e-5
            self.image_wight1 = self.epoch //5/4

            self.run_epoch()
            if (self.epoch + 1) % self.opt.save_frequency == 0:
                self.save_model()

    def run_epoch(self):
        """Run a single epoch of training and validation
        """
        #self.model_lr_scheduler.step()

        print("Training\n  ")
        save_log = os.path.join(self.log_path,'results1.txt')
        results_edit = open(save_log, mode='a')
        results_edit.write("Training")
        results_edit.close()
        self.set_train()

        for batch_idx, inputs in enumerate(self.train_loader):

            before_op_time = time.time()

            outputs, losses = self.process_batch(inputs)

            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)
                #pdb.set_trace()

                if "depth_gt" in inputs:
                    self.compute_depth_losses(inputs, outputs, losses)

                self.log("train", inputs, outputs, losses)
                self.val()

            self.step += 1
        self.model_lr_scheduler.step()

    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 = {}
            depth_fea = {}
            pose_fea = {}
            outputs={}
            for k in self.opt.frame_ids:
                features[k],depth_fea[k],pose_fea[k] = self.models["encoder"](inputs["color_aug", k, 0])
               # features[k]=features[k].detach()
                #pdb.set_trace()
                outputs.update(self.models["depth"](features[k],depth_fea[k],k))

        if self.opt.predictive_mask:
            outputs["predictive_mask"] = self.models["predictive_mask"](features)

        if self.use_pose_net:
            outputs.update(self.predict_poses(features,pose_fea))

        self.generate_pred(inputs, outputs)
        losses = self.compute_losses(inputs, outputs)
       # pdb.set_trace()

        return outputs, losses

    def predict_poses(self, features,pose_fea):
        """Predict poses between input frames for monocular sequences.
        """
        outputs = {}
        if self.num_pose_frames == 2:
            # In this setting, we compute the pose to each source frame via a
            # separate forward pass through the pose network.

            # select what features the pose network takes as input
            if self.opt.pose_model_type == "shared":
                pose_feats = {f_i: torch.cat([features[f_i][-1],pose_fea[f_i][-1]] ,1)for f_i in self.opt.frame_ids}
            else:
                # pose_feats = {f_i: inputs["color_aug", f_i, 0] for f_i in self.opt.frame_ids}
                pose_feats = {f_i: torch.cat([features[f_i][-1],pose_fea[f_i][-1] ],1)for f_i in self.opt.frame_ids}

            for f_i in self.opt.frame_ids[1:]:
                if f_i != "s":
                    # To maintain ordering we always pass frames in temporal order
                    if f_i < 0:
                        pose_inputs = [pose_feats[f_i], pose_feats[0]]
                    else:
                        pose_inputs = [pose_feats[0], pose_feats[f_i]]

                    # if self.opt.pose_model_type == "separate_resnet":
                    #     pose_inputs = [self.models["pose_encoder"](torch.cat(pose_inputs, 1))]
                    # elif self.opt.pose_model_type == "posecnn":
                    #     pose_inputs = torch.cat(pose_inputs, 1)

                    axisangle, translation = self.models["pose"](pose_inputs)
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation

                    # Invert the matrix if the frame id is negative
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0], invert=(f_i < 0))
                    outputs[("cam_T_cam", 0, f_i, "inv")] = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0], invert=(f_i > 0))

        else:
            # Here we input all frames to the pose net (and predict all poses) together
            if self.opt.pose_model_type in ["separate_resnet", "posecnn"]:
                pose_inputs = torch.cat(
                    [inputs[("color_aug", i, 0)] for i in self.opt.frame_ids if i != "s"], 1)

                if self.opt.pose_model_type == "separate_resnet":
                    pose_inputs = [self.models["pose_encoder"](pose_inputs)]

            elif self.opt.pose_model_type == "shared":
                pose_inputs = [features[i] for i in self.opt.frame_ids if i != "s"]

            axisangle, translation = self.models["pose"](pose_inputs)

            for i, f_i in enumerate(self.opt.frame_ids[1:]):
                if f_i != "s":
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, i], translation[:, i])

        return outputs

    def val(self):
        """Validate the model on a single minibatch
        """
        self.set_eval()
        try:
            inputs = self.val_iter.next()
        except StopIteration:
            self.val_iter = iter(self.val_loader)
            inputs = self.val_iter.next()

        with torch.no_grad():
            outputs, losses = self.process_batch(inputs)

            if "depth_gt" in inputs:
                self.compute_depth_losses(inputs, outputs, losses)

            self.log("val", inputs, outputs, losses)
            del inputs, outputs, losses

        self.set_train()

    def generate_pred(self, inputs, outputs):
        """Generate the warped (reprojected) color images for a minibatch.
        Generated images are saved into the `outputs` dictionary.
        """
        source_scale = 0
        for scale in self.opt.scales:
            for k in self.opt.frame_ids:
                disp = outputs[("disp",k, scale)]
                if self.opt.v1_multiscale:
                    source_scale = scale
                else:
                    disp = F.interpolate(
                        disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)

                _, depth = disp_to_depth(disp, self.opt.min_depth, self.opt.max_depth)

                outputs[("depth", k, scale)] = depth

            for i, frame_id in enumerate(self.opt.frame_ids[1:]):

                if frame_id == "s":
                    T = inputs["stereo_T"]
                    inv_T=inputs["stereo_inv_T"]
                else:
                    T = outputs[("cam_T_cam", 0, frame_id)]
                    inv_T = outputs[("cam_T_cam", 0, frame_id, "inv")]

                # from the authors of https://arxiv.org/abs/1712.00175
                if self.opt.pose_model_type == "posecnn":

                    axisangle = outputs[("axisangle", 0, frame_id)]
                    translation = outputs[("translation", 0, frame_id)]

                    inv_depth = 1 / outputs[("depth", 0, scale)]
                    mean_inv_depth = inv_depth.mean(3, True).mean(2, True)

                    T = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0] * mean_inv_depth[:, 0], frame_id < 0)

                cam_points = self.backproject_depth[source_scale](
                    outputs[("depth", 0, scale)], inputs[("inv_K", source_scale)])
                pix_coords,_ = self.project_3d[source_scale](
                    cam_points, inputs[("K", source_scale)], T)

                outputs[("sample", frame_id, scale)] = pix_coords

                outputs[("color", frame_id, scale)] = F.grid_sample(
                    inputs[("color", frame_id, source_scale)],
                    outputs[("sample", frame_id, scale)],
                    padding_mode="border",align_corners=True)
                if frame_id == "s":
                    pix_depth = outputs[("depth", frame_id, scale)]
                else:
                    cam_points_depth = self.backproject_depth[source_scale](
                    outputs[("depth", frame_id, scale)], inputs[("inv_K", source_scale)])

                    _, pix_depth = self.project_3d[source_scale](
                    cam_points_depth, inputs[("K", source_scale)], inv_T)

                outputs[("depth_inference", frame_id, scale)] = F.grid_sample(
                    pix_depth,
                    outputs[("sample", frame_id, scale)],
                    padding_mode="border", align_corners=True )

                if not self.opt.disable_automasking:
                    outputs[("color_identity", frame_id, scale)] = \
                        inputs[("color", frame_id, source_scale)]

    def compute_reprojection_loss(self, pred, target):
        """Computes reprojection loss between a batch of predicted and target images
        """
        abs_diff = torch.abs(target - pred)
        l1_loss = abs_diff.mean(1, True)

        if self.opt.no_ssim:
            reprojection_loss = l1_loss
        else:
            ssim_loss = self.ssim(pred, target).mean(1, True)
            reprojection_loss = 0.85 * ssim_loss + 0.15 * l1_loss

        return reprojection_loss

    @staticmethod
    def compute_loss_masks(reprojection_losses,identity_reprojection_losses,depth_reprojection_losses,identity_depth_losses,occlude_wight):
        """ Some codes will be made public after the paper is accepted."""


        return static_mask, occlude_mask,depth_static_mask

    @staticmethod
    def detect_outlier(depth_loss):
        # pdb.set_trace()
        mean = torch.mean(depth_loss, [1, 2, 3], keepdim=True)
       # mean = torch.quantile(depth_loss, 0.6, [1, 2, 3], keepdim=True,interpolation='midpoint')
        std = torch.std(depth_loss, [1, 2, 3], keepdim=True)
        mask_A = torch.abs(depth_loss-mean)/std
        #mask_B = (depth_loss - mean)/std > -3.8
        #mask = torch.logical_and(mask_A, mask_B)
        #
        return mask_A.float().detach()
    def compute_losses(self, inputs, outputs):
        """Compute the reprojection and smoothness losses for a minibatch
        """
        losses = {}
        total_loss = 0

        for scale in self.opt.scales:
            loss = 0
            to_optimise=0
            to_optimise_depth=0

            depth_mask=[]

            if self.opt.v1_multiscale:
                source_scale = scale
            else:
                source_scale = 0

            disp = outputs[("disp",0, scale)]
            color = inputs[("color", 0, scale)]
            target = inputs[("color", 0, source_scale)]
            target_depth = outputs[("depth", 0, source_scale)].detach()
            reprojection_losses = []
            depth_reprojection_losses = []
            identity_reprojection_losses = []
            identity_depth_losses = []
            occlude_mask =[]


            for frame_id in self.opt.frame_ids[1:]:

                pred = outputs[("color", frame_id, scale)]

                pred_depth = outputs[("depth_inference", frame_id, scale)].to(self.device)
                image_loss=self.compute_reprojection_loss(pred, target)
                reprojection_losses.append(image_loss)


                depth_losses=pred_depth - target_depth
                occlude=self.detect_outlier(depth_losses)
                
                        """Compute the reprojection and smoothness losses for a minibatch
        """


            to_optimise += combined.mean()
            to_optimise_depth += combined_depth.mean()


            loss += to_optimise
            loss += to_optimise_depth*self.opt.dep_wight
           # pdb.set_trace()


            mean_disp = disp.mean(2, True).mean(3, True)
            norm_disp = disp / (mean_disp + 1e-7)
            smooth_loss = get_smooth_loss(norm_disp, color)

            loss += self.opt.disparity_smoothness * smooth_loss / (2 ** scale)
            total_loss += loss
            losses["loss/{}".format(scale)] = loss

        total_loss /= self.num_scales
        losses["loss"] = total_loss
        return losses



    def log_time(self, batch_idx, duration, loss):
        """Print a logging statement to the terminal
        """
        samples_per_sec = self.opt.batch_size / duration
        time_sofar = time.time() - self.start_time
        training_time_left = (
                                     self.num_total_steps / self.step - 1.0) * time_sofar if self.step > 0 else 0
        print_string = "epoch {:>3} | batch {:>6} | examples/s: {:5.1f}" + \
                       " | loss: {:.5f} | time elapsed: {} | time left: {}"
        print(print_string.format(self.epoch, batch_idx, samples_per_sec, loss,
                                  sec_to_hm_str(time_sofar), sec_to_hm_str(training_time_left)))
        save_log = os.path.join(self.log_path,'results1.txt')
        results_edit = open(save_log, mode='a')
        results_edit.write("epoch {:>3} | batch {:>6} | examples/s: {:5.1f}| loss: {:.5f} | time elapsed: {} | time left: {}\n".format(self.epoch, batch_idx, samples_per_sec, loss,
                                  sec_to_hm_str(time_sofar), sec_to_hm_str(training_time_left)))
        results_edit.close()


    def log(self, mode, inputs, outputs, losses):
        """Write an event to the tensorboard events file
        """
        writer = self.writers[mode]
        for l, v in losses.items():
            writer.add_scalar("{}".format(l), v, self.step)

        for j in range(min(4, self.opt.batch_size)):  # write a maxmimum of four images
            for s in self.opt.scales:
                for frame_id in self.opt.frame_ids:
                    writer.add_image(
                        "color_{}_{}/{}".format(frame_id, s, j),
                        inputs[("color", frame_id, s)][j].data, self.step)

                    if s == 0 and frame_id != 0:
                        writer.add_image(
                            "color_pred_{}_{}/{}".format(frame_id, s, j),
                            outputs[("color", frame_id, s)][j].data, self.step)

                writer.add_image(
                    "disp_{}/{}".format(s, j),
                    normalize_image(outputs[("disp", 0,s)][j]), self.step)

                if self.opt.predictive_mask:
                    for f_idx, frame_id in enumerate(self.opt.frame_ids[1:]):
                        writer.add_image(
                            "predictive_mask_{}_{}/{}".format(frame_id, s, j),
                            outputs["predictive_mask"][("disp", s)][j, f_idx][None, ...],
                            self.step)

                # elif not self.opt.disable_automasking:
                #     writer.add_image(
                #         "automask_{}/{}".format(s, j),
                #         outputs["identity_selection/{}".format(s)][j][None, ...], self.step)
                if frame_id !=0:
                    writer.add_image(
                        "occlude_mask/{}_/{}/{}".format(frame_id,s,j),
                        outputs["occlude_mask/{}_/{}".format(frame_id,s)][j].data, self.step)

    def save_opts(self):
        """Save options to disk so we know what we ran this experiment with
        """
        models_dir = os.path.join(self.log_path, "models")
        if not os.path.exists(models_dir):
            os.makedirs(models_dir)
        to_save = self.opt.__dict__.copy()

        with open(os.path.join(models_dir, 'opt.json'), 'w') as f:
            json.dump(to_save, f, indent=2)

    def save_model(self):
        """Save model weights to disk
        """
        save_folder = os.path.join(self.log_path, "models", "weights_{}".format(self.epoch))
        if not os.path.exists(save_folder):
            os.makedirs(save_folder)

        for model_name, model in self.models.items():
            save_path = os.path.join(save_folder, "{}.pth".format(model_name))
            to_save = model.module.state_dict()
            if model_name == 'encoder':
                # save the sizes - these are needed at prediction time
                to_save['height'] = self.opt.height
                to_save['width'] = self.opt.width
                to_save['use_stereo'] = self.opt.use_stereo
            torch.save(to_save, save_path)

        save_path = os.path.join(save_folder, "{}.pth".format("adam"))
        torch.save(self.model_optimizer.state_dict(), save_path)

    def load_model(self):
        """Load model(s) from disk
        """
        self.opt.load_weights_folder = os.path.expanduser(self.opt.load_weights_folder)

        assert os.path.isdir(self.opt.load_weights_folder), \
            "Cannot find folder {}".format(self.opt.load_weights_folder)
        print("loading model from folder {}".format(self.opt.load_weights_folder))
        if not os.path.exists(self.log_path):
            os.makedirs(self.log_path)
        save_log = os.path.join(self.log_path,'results1.txt')
        results_edit = open(save_log, mode='a')
        results_edit.write("loading model from folder {}".format(self.opt.load_weights_folder))
        results_edit.close()

        for n in self.opt.models_to_load:
            print("Loading {} weights...".format(n))
            save_log = os.path.join(self.log_path, 'results1.txt')
            results_edit = open(save_log, mode='a')
            results_edit.write("Loading {} weights...".format(n))
            results_edit.close()
            path = os.path.join(self.opt.load_weights_folder, "{}.pth".format(n))
            model_dict = self.models[n].module.state_dict()
            pretrained_dict = torch.load(path)
            if n=="pose":
                # key_map = {
                #     "net.0.weight": "squeeze.weight",
                #     "net.0.bias": "squeeze.bias",
                #     "net.1.weight": "pose_0.weight",
                #     "net.1.bias": "pose_0.bias",
                #     "net.2.weight": "pose_1.weight",
                #     "net.2.bias": "pose_1.bias",
                #     "net.3.weight": "pose_2.weight",
                #     "net.3.bias": "pose_2.bias",
                # }
                # new_state_dict = {}
                # for key, value in pretrained_dict.items():
                #     if key in key_map:
                #         new_key = key_map[key]
                #         new_state_dict[new_key] = value
                # self.models[n].module.load_state_dict(new_state_dict)
                pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
                model_dict.update(pretrained_dict)
                self.models[n].module.load_state_dict(model_dict)
                print("load posedconder")
                save_log = os.path.join(self.log_path, 'results1.txt')
                results_edit = open(save_log, mode='a')
                results_edit.write("load posedconder")
                results_edit.close()
            else:
                pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
                model_dict.update(pretrained_dict)
                self.models[n].module.load_state_dict(model_dict)

        # loading adam state
        optimizer_load_path = os.path.join(self.opt.load_weights_folder, "adam.pth")
        if os.path.isfile(optimizer_load_path):
            print("Loading Adam weights")
            save_log = os.path.join(self.log_path, 'results1.txt')
            results_edit = open(save_log, mode='a')
            results_edit.write("Loading Adam weights")
            results_edit.close()

            optimizer_dict = torch.load(optimizer_load_path)
            #self.model_optimizer.load_state_dict(optimizer_dict)
        else:
            print("Cannot find Adam weights so Adam is randomly initialized")
            save_log = os.path.join(self.log_path, 'results1.txt')
            results_edit = open(save_log, mode='a')
            results_edit.write("Cannot find Adam weights so Adam is randomly initialized")
            results_edit.close()