(self, args)
| 17 | |
| 18 | class Model: |
| 19 | def __init__(self, args): |
| 20 | |
| 21 | if args.vis: |
| 22 | self.args = args |
| 23 | return |
| 24 | |
| 25 | cudnn.benchmark = True |
| 26 | |
| 27 | init_distributed_mode(args) |
| 28 | |
| 29 | self.args = args |
| 30 | self.device = torch.device("cpu" if self.args.no_cuda or not torch.cuda.is_available() else "cuda") |
| 31 | |
| 32 | self.network = MVSNet(ndepths=args.ndepths, depth_interval_ratio=args.interval_ratio, fea_mode=args.fea_mode, |
| 33 | agg_mode=args.agg_mode, depth_mode=args.depth_mode, |
| 34 | winner_take_all_to_generate_depth=args.winner_take_all_to_generate_depth,inverse_depth=self.args.inverse_depth).to(self.device) |
| 35 | |
| 36 | if self.args.distributed and self.args.sync_bn: |
| 37 | self.network = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.network) |
| 38 | |
| 39 | if not (self.args.val or self.args.test): |
| 40 | |
| 41 | self.optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.network.parameters()), lr=args.lr, |
| 42 | weight_decay=args.wd) |
| 43 | self.lr_scheduler = get_schedular(self.optimizer, self.args) |
| 44 | self.train_loader, self.train_sampler = get_loader(args, args.datapath, args.trainlist, args.nviews, "train") |
| 45 | |
| 46 | if not self.args.test: |
| 47 | self.loss_func = mvs_loss |
| 48 | |
| 49 | self.val_loader, self.val_sampler = get_loader(args, args.datapath, args.testlist, 5, "test",force_test=True) |
| 50 | if is_main_process(): |
| 51 | self.writer = SummaryWriter(log_dir=args.log_dir, comment="Record network info") |
| 52 | |
| 53 | self.network_without_ddp = self.network |
| 54 | if self.args.distributed: |
| 55 | self.network = DistributedDataParallel(self.network, device_ids=[self.args.local_rank]) |
| 56 | # self.network = DistributedDataParallel(self.network, device_ids=[self.args.local_rank],find_unused_parameters=True) |
| 57 | self.network_without_ddp = self.network.module |
| 58 | |
| 59 | if self.args.resume: |
| 60 | checkpoint = torch.load(self.args.resume, map_location="cpu") |
| 61 | if not (self.args.val or self.args.test or self.args.blendedmvs_finetune): |
| 62 | self.args.start_epoch = checkpoint["epoch"] + 1 |
| 63 | self.optimizer.load_state_dict(checkpoint["optimizer"]) |
| 64 | self.lr_scheduler.load_state_dict(checkpoint["lr_scheduler"]) |
| 65 | import collections |
| 66 | new_dic=collections.OrderedDict() |
| 67 | for (key,values) in checkpoint["model"].items(): |
| 68 | if "attn_mask" not in key: |
| 69 | new_dic[key]=values |
| 70 | self.network_without_ddp.load_state_dict(new_dic) |
| 71 | |
| 72 | self.blendmvs=('dataset_low_res' in args.datapath) |
| 73 | |
| 74 | def main(self): |
| 75 | # print(self.args.test) |
nothing calls this directly
no test coverage detected