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)