Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
model = UNet(n_channels=3, n_classes=num_classes)#UNet2(3,num_classes)
elif architecture == 'fast_scnn':
model = get_fast_scnn(num_classes)
elif architecture == 'nested_unet':
print('Nested UNET is deprecated for now, defaulting to UNET.')
model = UNet(n_channels=3, n_classes=num_classes)#NestedUNet(3, num_classes)
elif architecture.startswith('efficientnet'):
from efficientnet_pytorch import EfficientNet
if pretrain:
model = EfficientNet.from_pretrained(architecture, override_params=dict(num_classes=num_classes))
else:
model = EfficientNet.from_name(architecture, override_params=dict(num_classes=num_classes))
print(model)
elif architecture.startswith('sqnxt'):
from pytorchcv.model_provider import get_model as ptcv_get_model
model = ptcv_get_model(architecture, pretrained=pretrain)
num_ftrs=int(128*int(architecture.split('_')[-1][1]))
model.output=MLP(num_ftrs, [1000], dropout_p=0., n_outputs=num_classes, binary=add_sigmoid, softmax=False).mlp
else:
#for pretrained on imagenet
model_names = [m for m in dir(models) if not m.startswith('__')]
segmentation_model_names = [m for m in dir(segmodels) if not m.startswith('__')]
if architecture in model_names:
model = getattr(models, architecture)(pretrained=pretrain)
if segmentation:
if architecture in segmentation_model_names:
model = getattr(segmodels, architecture)(pretrained=pretrain)
else:
model = UNet(n_channels=3, n_classes=num_classes)
if architecture.startswith('deeplab'):
model.classifier[4] = nn.Conv2d(256, num_classes, kernel_size=(1, 1), stride=(1, 1))
model = FixedSegmentationModule(model)