def test_custom_csp():
with pytest.raises(AssertionError):
cfg = dict(
name="CustomCspNet", net_cfg=[["Conv", 3, 32, 3, 2]], out_stages=(8, 9)
)
build_backbone(cfg)
with pytest.raises(AssertionError):
CustomCspNet(net_cfg=dict(a=1), out_stages=(0, 1), activation="ReLU6")
input = torch.rand(1, 3, 64, 64)
out_stages = (0, 1, 2, 3, 4, 5)
net_cfg = [
["Conv", 3, 32, 3, 2], # 1/2
["MaxPool", 3, 2], # 1/4
["CspBlock", 32, 1, 3, 1], # 1/4
["CspBlock", 64, 2, 3, 2], # 1/8
["CspBlock", 128, 2, 3, 2], # 1/16
["CspBlock", 256, 3, 3, 2], # 1/32
]
model = CustomCspNet(net_cfg=net_cfg, out_stages=out_stages, activation="ReLU6")
output = model(input)
assert output[0].shape == (1, 32, 32, 32)
assert output[1].shape == (1, 32, 16, 16)
assert output[2].shape == (1, 64, 16, 16)
assert output[3].shape == (1, 128, 8, 8)
assert output[4].shape == (1, 256, 4, 4)
assert output[5].shape == (1, 512, 2, 2)
def test_repvgg():
with pytest.raises(AssertionError):
cfg = dict(name="RepVGG", arch="A3")
build_backbone(cfg)
with pytest.raises(AssertionError):
RepVGG(arch="A0", out_stages=(4, 5, 6))
input = torch.rand(1, 3, 64, 64)
model = RepVGG(arch="A0", out_stages=(1, 2, 3, 4), activation="PReLU")
output = model(input)
assert output[0].shape == (1, 48, 16, 16)
assert output[1].shape == (1, 96, 8, 8)
assert output[2].shape == (1, 192, 4, 4)
assert output[3].shape == (1, 1280, 2, 2)
# test last channel
model = RepVGG(arch="A1", out_stages=(1, 2, 3, 4), last_channel=640)
output = model(input)
assert output[0].shape == (1, 64, 16, 16)
assert output[1].shape == (1, 128, 8, 8)
assert output[2].shape == (1, 256, 4, 4)
assert output[3].shape == (1, 640, 2, 2)
deploy_model = RepVGG(arch="A0", deploy=True)
deploy_model = repvgg_model_convert(model, deploy_model, save_path=None)
dep_output = deploy_model(input)
assert dep_output[0].shape == (1, 64, 16, 16)
assert dep_output[1].shape == (1, 128, 8, 8)
assert dep_output[2].shape == (1, 256, 4, 4)
assert dep_output[3].shape == (1, 640, 2, 2)
def test_mobilenetv2():
with pytest.raises(AssertionError):
cfg = dict(name="MobileNetV2", width_mult=1.0, out_stages=(8, 9))
build_backbone(cfg)
input = torch.rand(1, 3, 64, 64)
out_stages = (0, 1, 2, 3, 4, 5, 6)
model = MobileNetV2(width_mult=1.0,
out_stages=out_stages,
activation="ReLU6")
output = model(input)
assert output[0].shape == (1, 16, 32, 32)
assert output[1].shape == (1, 24, 16, 16)
assert output[2].shape == (1, 32, 8, 8)
assert output[3].shape == (1, 64, 4, 4)
assert output[4].shape == (1, 96, 4, 4)
assert output[5].shape == (1, 160, 2, 2)
assert output[6].shape == (1, 1280, 2, 2)
model = MobileNetV2(width_mult=0.75,
out_stages=out_stages,
activation="LeakyReLU")
output = model(input)
assert output[0].shape == (1, 12, 32, 32)
assert output[1].shape == (1, 18, 16, 16)
assert output[2].shape == (1, 24, 8, 8)
assert output[3].shape == (1, 48, 4, 4)
assert output[4].shape == (1, 72, 4, 4)
assert output[5].shape == (1, 120, 2, 2)
assert output[6].shape == (1, 1280, 2, 2)
def test_shufflenetv2():
with pytest.raises(NotImplementedError):
cfg = dict(name="ShuffleNetV2", model_size="3.0x", pretrain=False)
build_backbone(cfg)
with pytest.raises(AssertionError):
ShuffleNetV2("1.0x", out_stages=(1, 2, 3))
input = torch.rand(1, 3, 64, 64)
model = ShuffleNetV2(model_size="0.5x",
out_stages=(2, 3, 4),
pretrain=True)
output = model(input)
assert output[0].shape == (1, 48, 8, 8)
assert output[1].shape == (1, 96, 4, 4)
assert output[2].shape == (1, 192, 2, 2)
model = ShuffleNetV2(model_size="0.5x", out_stages=(3, 4), pretrain=True)
output = model(input)
assert output[0].shape == (1, 96, 4, 4)
assert output[1].shape == (1, 192, 2, 2)
model = ShuffleNetV2(model_size="1.0x",
pretrain=False,
with_last_conv=True)
assert hasattr(model.stage4, "conv5")
output = model(input)
assert output[0].shape == (1, 116, 8, 8)
assert output[1].shape == (1, 232, 4, 4)
assert output[2].shape == (1, 1024, 2, 2)
model = ShuffleNetV2(model_size="1.5x",
pretrain=False,
with_last_conv=False,
activation="ReLU6")
assert not hasattr(model.stage4, "conv5")
output = model(input)
assert output[0].shape == (1, 176, 8, 8)
assert output[1].shape == (1, 352, 4, 4)
assert output[2].shape == (1, 704, 2, 2)
model = ShuffleNetV2(model_size="2.0x",
pretrain=False,
with_last_conv=False)
output = model(input)
assert output[0].shape == (1, 244, 8, 8)
assert output[1].shape == (1, 488, 4, 4)
assert output[2].shape == (1, 976, 2, 2)
def test_timm_wrapper():
cfg = dict(
name="TIMMWrapper",
model_name="resnet18",
features_only=True,
pretrained=False,
output_stride=32,
out_indices=(1, 2, 3, 4),
)
model = build_backbone(cfg)
input = torch.rand(1, 3, 64, 64)
output = model(input)
assert len(output) == 4
assert output[0].shape == (1, 64, 16, 16)
assert output[1].shape == (1, 128, 8, 8)
assert output[2].shape == (1, 256, 4, 4)
assert output[3].shape == (1, 512, 2, 2)
model = TIMMWrapper(
model_name="mobilenetv3_large_100",
features_only=True,
pretrained=False,
output_stride=32,
out_indices=(1, 2, 3, 4),
)
output = model(input)
assert len(output) == 4
assert output[0].shape == (1, 24, 16, 16)
assert output[1].shape == (1, 40, 8, 8)
assert output[2].shape == (1, 112, 4, 4)
assert output[3].shape == (1, 960, 2, 2)
def test_resnet():
with pytest.raises(KeyError):
cfg = dict(name="ResNet", depth=15)
build_backbone(cfg)
with pytest.raises(AssertionError):
ResNet(depth=18, out_stages=(4, 5, 6))
input = torch.rand(1, 3, 64, 64)
model = ResNet(depth=18,
out_stages=(1, 2, 3, 4),
activation="PReLU",
pretrain=True)
output = model(input)
assert output[0].shape == (1, 64, 16, 16)
assert output[1].shape == (1, 128, 8, 8)
assert output[2].shape == (1, 256, 4, 4)
assert output[3].shape == (1, 512, 2, 2)
model = ResNet(depth=34,
out_stages=(1, 2, 3, 4),
activation="LeakyReLU",
pretrain=False)
output = model(input)
assert output[0].shape == (1, 64, 16, 16)
assert output[1].shape == (1, 128, 8, 8)
assert output[2].shape == (1, 256, 4, 4)
assert output[3].shape == (1, 512, 2, 2)
model = ResNet(depth=50, out_stages=(1, 2, 3, 4), pretrain=False)
output = model(input)
assert output[0].shape == (1, 256, 16, 16)
assert output[1].shape == (1, 512, 8, 8)
assert output[2].shape == (1, 1024, 4, 4)
assert output[3].shape == (1, 2048, 2, 2)
model = ResNet(depth=101, out_stages=(1, 2, 3, 4), pretrain=False)
output = model(input)
assert output[0].shape == (1, 256, 16, 16)
assert output[1].shape == (1, 512, 8, 8)
assert output[2].shape == (1, 1024, 4, 4)
assert output[3].shape == (1, 2048, 2, 2)
def test_ghostnet():
with pytest.raises(AssertionError):
cfg = dict(name="GhostNet",
width_mult=1.0,
out_stages=(11, 12),
pretrain=False)
build_backbone(cfg)
input = torch.rand(1, 3, 64, 64)
out_stages = [i for i in range(10)]
model = GhostNet(width_mult=1.0,
out_stages=out_stages,
activation="ReLU6",
pretrain=True)
output = model(input)
assert output[0].shape == torch.Size([1, 16, 32, 32])
assert output[1].shape == torch.Size([1, 24, 16, 16])
assert output[2].shape == torch.Size([1, 24, 16, 16])
assert output[3].shape == torch.Size([1, 40, 8, 8])
assert output[4].shape == torch.Size([1, 40, 8, 8])
assert output[5].shape == torch.Size([1, 80, 4, 4])
assert output[6].shape == torch.Size([1, 112, 4, 4])
assert output[7].shape == torch.Size([1, 160, 2, 2])
assert output[8].shape == torch.Size([1, 160, 2, 2])
assert output[9].shape == torch.Size([1, 960, 2, 2])
model = GhostNet(width_mult=0.75,
out_stages=out_stages,
activation="LeakyReLU",
pretrain=False)
output = model(input)
assert output[0].shape == torch.Size([1, 12, 32, 32])
assert output[1].shape == torch.Size([1, 20, 16, 16])
assert output[2].shape == torch.Size([1, 20, 16, 16])
assert output[3].shape == torch.Size([1, 32, 8, 8])
assert output[4].shape == torch.Size([1, 32, 8, 8])
assert output[5].shape == torch.Size([1, 60, 4, 4])
assert output[6].shape == torch.Size([1, 84, 4, 4])
assert output[7].shape == torch.Size([1, 120, 2, 2])
assert output[8].shape == torch.Size([1, 120, 2, 2])
assert output[9].shape == torch.Size([1, 720, 2, 2])
def test_efficientnet_lite():
with pytest.raises(AssertionError):
cfg = dict(
name="EfficientNetLite",
model_name="efficientnet_lite0",
out_stages=(7, 8, 9),
)
build_backbone(cfg)
with pytest.raises(AssertionError):
EfficientNetLite(model_name="efficientnet_lite9")
input = torch.rand(1, 3, 64, 64)
model = EfficientNetLite(model_name="efficientnet_lite0",
out_stages=(0, 1, 2, 3, 4, 5, 6),
pretrain=True)
output = model(input)
assert output[0].shape == (1, 16, 32, 32)
assert output[1].shape == (1, 24, 16, 16)
assert output[2].shape == (1, 40, 8, 8)
assert output[3].shape == (1, 80, 4, 4)
assert output[4].shape == (1, 112, 4, 4)
assert output[5].shape == (1, 192, 2, 2)
assert output[6].shape == (1, 320, 2, 2)
model = EfficientNetLite(
model_name="efficientnet_lite1",
out_stages=(0, 1, 2, 3, 4, 5, 6),
pretrain=False,
)
output = model(input)
assert output[0].shape == (1, 16, 32, 32)
assert output[1].shape == (1, 24, 16, 16)
assert output[2].shape == (1, 40, 8, 8)
assert output[3].shape == (1, 80, 4, 4)
assert output[4].shape == (1, 112, 4, 4)
assert output[5].shape == (1, 192, 2, 2)
assert output[6].shape == (1, 320, 2, 2)
model = EfficientNetLite(
model_name="efficientnet_lite2",
out_stages=(0, 1, 2, 3, 4, 5, 6),
activation="ReLU",
pretrain=False,
)
output = model(input)
assert output[0].shape == (1, 16, 32, 32)
assert output[1].shape == (1, 24, 16, 16)
assert output[2].shape == (1, 48, 8, 8)
assert output[3].shape == (1, 88, 4, 4)
assert output[4].shape == (1, 120, 4, 4)
assert output[5].shape == (1, 208, 2, 2)
assert output[6].shape == (1, 352, 2, 2)
model = EfficientNetLite(
model_name="efficientnet_lite3",
out_stages=(0, 1, 2, 3, 4, 5, 6),
pretrain=False,
)
output = model(input)
assert output[0].shape == (1, 24, 32, 32)
assert output[1].shape == (1, 32, 16, 16)
assert output[2].shape == (1, 48, 8, 8)
assert output[3].shape == (1, 96, 4, 4)
assert output[4].shape == (1, 136, 4, 4)
assert output[5].shape == (1, 232, 2, 2)
assert output[6].shape == (1, 384, 2, 2)
model = EfficientNetLite(
model_name="efficientnet_lite4",
out_stages=(0, 1, 2, 3, 4, 5, 6),
pretrain=False,
)
output = model(input)
assert output[0].shape == (1, 24, 32, 32)
assert output[1].shape == (1, 32, 16, 16)
assert output[2].shape == (1, 56, 8, 8)
assert output[3].shape == (1, 112, 4, 4)
assert output[4].shape == (1, 160, 4, 4)
assert output[5].shape == (1, 272, 2, 2)
assert output[6].shape == (1, 448, 2, 2)