def test_build_fpn_with_repvgg(self): r"""Test building a 3D FPN model with RepVGG backbone from configs. """ cfg = get_cfg_defaults() cfg.MODEL.ARCHITECTURE = 'fpn_3d' cfg.MODEL.BACKBONE = 'repvgg' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = build_model(cfg, device).eval() message = "Get unexpected model architecture!" arch_name = model.module.__class__.__name__ self.assertEqual(arch_name, "FPN3D", message) message = "No RepVGG block in the backbone!" count = 0 for layer in model.modules(): if isinstance(layer, RepVGGBlock3D): count += 1 self.assertGreater(count, 0) # test the weight conversion when using RepVGG as backbone model.eval() train_dict = model.module.state_dict() deploy_dict = RepVGG3D.repvgg_convert_as_backbone(train_dict) cfg.MODEL.DEPLOY_MODE = True deploy_model = build_model(cfg, device).eval() deploy_model.module.load_state_dict(deploy_dict, strict=True) x = torch.rand(2, 1, 9, 65, 65) y1 = model(x) y2 = deploy_model(x) self.assertTrue(torch.allclose(y1, y2, atol=1e-4))
def main(): r"""Main function. """ # arguments args = get_args() print("Command line arguments:") print(args) # configurations cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) if args.inference: update_inference_config(cfg) cfg.freeze() print("Configuration details:") print(cfg) if not os.path.exists(cfg.DATASET.OUTPUT_PATH): print('Output directory: ', cfg.DATASET.OUTPUT_PATH) os.makedirs(cfg.DATASET.OUTPUT_PATH) save_all_cfg(cfg, cfg.DATASET.OUTPUT_PATH) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Device: ", device) mode = 'test' if args.inference else 'train' trainer = Trainer(cfg, device, mode, args.checkpoint) if args.inference: trainer.test() else: trainer.train()
def test_build_default_model(self): r"""Test building model from configs. """ cfg = get_cfg_defaults() cfg.freeze() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = build_model(cfg, device) self.assertTrue( isinstance(model, (torch.nn.Module, torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)))
def test_build_fpn_with_botnet(self): r"""Test building a 3D FPN model with BotNet3D backbone from configs. """ cfg = get_cfg_defaults() cfg.MODEL.ARCHITECTURE = 'fpn_3d' cfg.MODEL.BACKBONE = 'botnet' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = build_model(cfg, device).eval() d, h, w = cfg.MODEL.INPUT_SIZE x = torch.rand(2, 1, d, h, w) y1 = model(x) self.assertTupleEqual(tuple(y1.shape), (2, 1, d, h, w))
def test_build_deeplab_with_resnet(self): r"""Test building 2D deeplabv3 model with resnet backbone. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") cfg = get_cfg_defaults() cfg.MODEL.BACKBONE = 'resnet101' cfg.MODEL.AUX_OUT = True c_i = cfg.MODEL.IN_PLANES c_o = cfg.MODEL.OUT_PLANES b, h, w = 2, 65, 65 x = torch.rand(b, c_i, h, w).to(device) for arch in ['deeplabv3a', 'deeplabv3b', 'deeplabv3c']: cfg.MODEL.ARCHITECTURE = arch model = build_model(cfg, device).eval() y = model(x) self.assertTrue(isinstance(y, OrderedDict)) self.assertTrue("aux" in y.keys()) for key in y.keys(): self.assertTupleEqual(tuple(y[key].shape), (b, c_o, h, w))
def test_build_fpn_with_efficientnet(self): r"""Test building a 3D FPN model with EfficientNet3D backbone from configs. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") cfg = get_cfg_defaults() cfg.MODEL.ARCHITECTURE = 'fpn_3d' cfg.MODEL.BACKBONE = 'efficientnet' d, h, w = 9, 65, 65 x = torch.rand(2, 1, d, h, w).to(device) # inverted residual blocks cfg.MODEL.BLOCK_TYPE = 'inverted_res' model = build_model(cfg, device).eval() y1 = model(x) self.assertTupleEqual(tuple(y1.shape), (2, 1, d, h, w)) # inverted residual blocks with dilation cfg.MODEL.BLOCK_TYPE = 'inverted_res_dilated' model = build_model(cfg, device).eval() y1 = model(x) self.assertTupleEqual(tuple(y1.shape), (2, 1, d, h, w))
def main(): r"""Main function. """ # arguments args = get_args() print("Command line arguments:") print(args) # configurations cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) if args.inference: update_inference_cfg(cfg) cfg.freeze() print("Configuration details:") print(cfg) if not os.path.exists(cfg.dataset.output_path): print('Output directory: ', cfg.dataset.output_path) os.makedirs(cfg.dataset.output_path) save_all_cfg(cfg, cfg.dataset.output_path) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Device: ", device) cudnn.enabled = True cudnn.benchmark = True mode = 'test' if args.inference else 'train' trainer = Trainer(cfg, device, mode, args.checkpoint) if cfg.dataset.DO_CHUNK_TITLE == 0: if args.inference: trainer.test() else: trainer.train() else: trainer.run_chunk(mode)