# ============================================================================== # Copyright (c) 2022 The PersFormer Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import os import random import logging import subprocess import numpy as np logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s',) import torch import torch.nn as nn import torch.distributed as dist import torch.multiprocessing as mp from torch.nn import DataParallel as DP from torch.utils.data import Dataset, DataLoader from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler as DS from core.raft_stereo import RAFTStereo from core.raft_stereo_disp import RAFTStereoDisp from core.raft_stereo_mast3r import RAFTStereoMast3r from core.raft_stereo_depthany import RAFTStereoDepthAny from core.raft_stereo_noctx import RAFTStereoNoCTX from core.raft_stereo_depthfusion import RAFTStereoDepthFusion from core.raft_stereo_depthbeta import RAFTStereoDepthBeta from core.raft_stereo_depthbeta_nolbp import RAFTStereoDepthBetaNoLBP from core.raft_stereo_depthmatch import RAFTStereoDepthMatch from core.raft_stereo_depthbeta_refine import RAFTStereoDepthBetaRefine from core.raft_stereo_depth_postfusion import RAFTStereoDepthPostFusion from core.raft_stereo_metric3d import RAFTStereoMetric3D def setup_distributed(args): args.rank = int(os.getenv("RANK")) args.local_rank = int(os.getenv("LOCAL_RANK")) args.world_size = int(os.getenv("WORLD_SIZE")) # print("-"*10, "local_rank: {}, world_size:{}".format(args.local_rank, args.world_size), # " - {}, {}".format(dist.get_rank(), dist.get_world_size())) # they result in the same value dist.init_process_group(backend='nccl') torch.cuda.set_device(args.local_rank) torch.set_printoptions(precision=10) def ddp_init(args): if 'WORLD_SIZE' in os.environ: args.distributed = int(os.environ['WORLD_SIZE']) >= 1 if args.distributed: setup_distributed(args) # deterministic torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True torch.manual_seed(args.local_rank) np.random.seed(args.local_rank) random.seed(args.local_rank) print("complete initialization of local_rank:{}".format(args.local_rank)) def ddp_close(): dist.destroy_process_group() def to_python_float(t): if hasattr(t, 'item'): return t.item() else: return t[0] def reduce_tensor(tensor, world_size): rt = tensor.clone() dist.all_reduce(rt, op=dist.ReduceOp.SUM) rt /= world_size return rt def reduce_tensors(*tensors, world_size): return [reduce_tensor(tensor, world_size) for tensor in tensors] def get_loader(dataset, args): """ create dataset from ground-truth return a batch sampler based ont the dataset """ if args.distributed: if args.local_rank == 0: print('use distributed sampler') data_sampler = DS(dataset, shuffle=True, drop_last=True) data_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=data_sampler, num_workers=args.num_workers, pin_memory=True, persistent_workers=True) else: if args.local_rank == 0: print("use default sampler") # data_sampler = torch.utils.data.sampler.SubsetRandomSampler(sample_idx) # data_loader = DataLoader(transformed_dataset, # batch_size=args.batch_size, sampler=data_sampler, # num_workers=args.nworkers, pin_memory=True, # persistent_workers=True, # worker_init_fn=seed_worker, # generator=g) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, pin_memory=True, shuffle=True, num_workers=args.num_workers, drop_last=True) return data_loader NODE_RANK = os.getenv('NODE_RANK', default=0) LOCAL_RANK = os.getenv("LOCAL_RANK", default=0) def get_model_ddp(args): if args.model_name.lower() == "raftstereo": model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereo(args)) elif args.model_name.lower() == "raftstereodisp": model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDisp(args)) elif args.model_name.lower() == "raftstereomast3r": model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoMast3r(args)) elif args.model_name.lower() == "raftstereodepthany": model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthAny(args)) elif args.model_name.lower() == "raftstereodepthfusion": model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthFusion(args)) elif args.model_name.lower() == "RAFTStereoDepthBeta".lower(): model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthBeta(args)) elif args.model_name.lower() == "RAFTStereoDepthBetaNoLBP".lower(): model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthBetaNoLBP(args)) elif args.model_name.lower() == "RAFTStereoDepthMatch".lower(): model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthMatch(args)) elif args.model_name.lower() == "RAFTStereoDepthBetaRefine".lower(): model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthBetaRefine(args)) elif args.model_name.lower() == "RAFTStereoDepthPostFusion".lower(): model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoDepthPostFusion(args)) elif args.model_name.lower() == "RAFTStereoMetric3D".lower(): model = nn.SyncBatchNorm.convert_sync_batchnorm(RAFTStereoMetric3D(args)) else : raise Exception("No such model: {}".format(args.model_name)) device = torch.device("cuda", args.local_rank) model = model.to(device) if args.restore_ckpt is not None: assert args.restore_ckpt.endswith(".pth") or args.restore_ckpt.endswith(".tar") if args.local_rank==0 : logging.info("Loading checkpoint from {} ...".format(args.restore_ckpt)) checkpoint = torch.load(args.restore_ckpt) new_state_dict = {} for key, value in checkpoint.items(): new_key = key.replace('module.', '') # 去掉 'module.' 前缀 # if key.find("refinement.conf_estimate") != -1: # continue new_state_dict[new_key] = value # model.load_state_dict(new_state_dict, strict=True) model.load_state_dict(new_state_dict, strict=False) if args.local_rank==0 : logging.info(f"Done loading checkpoint") dist.barrier() # DDP setting if args.distributed: model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) else: model = DP(RAFTStereo(args)) return model