def test_UNet():
x = torch.rand((1, 1, 32 * 12, 32 * 12))
with torch.no_grad():
# unet
unet = UNet()
unet(x)
# custom encoder
unet = UNet(
encoder=lambda *args, **kwargs: ResNet.resnet26(*args, **kwargs).encoder,
)
unet(x)
# change decoder
unet = UNet(decoder=partial(UNetDecoder, widths=[256, 128, 64, 32, 16]))
unet(x)
# using efficienet net
unet = UNet(
encoder=lambda *args, **kwargs: EfficientNet.efficientnet_b2(
*args, **kwargs
).encoder
)
unet(x)
# combine them
unet = UNet(
encoder=lambda *args, **kwargs: EfficientNet.efficientnet_b2(
*args, **kwargs
).encoder,
decoder=partial(UNetDecoder, widths=[256, 128, 64, 32, 16]),
)
unet(x)
unet = UNet(
encoder=lambda *args, **kwargs: EfficientNetLite.efficientnet_lite3(
*args, **kwargs
).encoder,
)
unet(x)
# customize the encoder
unet = UNet(
encoder=partial(ResNetEncoder, block=ResNetBasicBlock, depths=[1, 1, 2, 2])
)
unet(x)
unet = UNet(
encoder=partial(
ResNetEncoder, block=ResNetBottleneckBlock, depths=[1, 1, 2, 2]
)
)
unet(x)
# custom block
unet = UNet(encoder=partial(UNetEncoder, block=SENetBasicBlock))
unet(x)
# using .from_encoder
unet = UNet.from_encoder(
lambda *args, **kwargs: ResNet.resnet26(*args, **kwargs)
)
unet(x)
# with AutoModel
unet = UNet.from_encoder(partial(AutoModel.from_name, "resnet18"))
unet(x)
with pytest.raises(AttributeError):
unet = UNet.from_encoder(lambda *args, **kwargs: None)
def test_EfficientNetLite2():
x = torch.rand(1, 3, 224, 224)
model = EfficientNetLite.efficientnet_lite2()
pred = model(x)
assert pred.shape[-1] == 1000