How to use the pytorchcv.model_provider.get_model function in pytorchcv

To help you get started, we’ve selected a few pytorchcv examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github jlevy44 / PathFlowAI / pathflowai / models.py View on Github external
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)

pytorchcv

Computer vision models for PyTorch

MIT
Latest version published 3 months ago

Package Health Score

61 / 100
Full package analysis

Popular pytorchcv functions