Neural network models is what deep learning is all about! While you can download some standard models from torchvision, we strive to create a library of models that are on the cutting edge of AI. Whenever possible, we provide pretrained solutions as well!

That said, we didn’t come up with any of these on our own so we owe a huge debt of gratitude to the many researchers who have shared their models and weights on github.

Caution: While we strive to ensure that all models can be used out of the box, sometimes things become broken due to Pytorch updates or misalignment of the planets. Please don’t yell at us. Gently point out what’s broken, or even better, submit a pull request to fix it!

Here Be Dragons: Aaand one more thing - we constantly plumb the depths of github for new models or tweaks to existing ones. While we don’t list this in the docs, there is a special testnets directory with tons of probably broken, semi-working, and at times crazy awesome models and model-variations. If you’re interested in the bleeding edge, that’s where you’d look (see for what’s available)

utility functions

pywick.models.model_utils.load_checkpoint(checkpoint_path, model=None, device='cpu', strict=True, ignore_chkpt_layers=None)[source]

Loads weights from a checkpoint into memory. If model is not None then the weights are loaded into the model.

  • checkpoint_path – (string): path to a pretrained network to load weights from
  • model – the model object to load weights onto (default: None)
  • device – (string): which device to load model onto (default:’cpu’)
  • strict – (bool): whether to ensure strict key matching (True) or to ignore non-matching keys. (default: True)
  • ignore_chkpt_layers – one of {string, list) – CURRENTLY UNIMPLEMENTED: whether to ignore some subset of layers from checkpoint. This is usually done when loading checkpoint data into a model with a different number of final classes. In that case, you can pass in a special string: ‘last_layer’ which will trigger the logic to chop off the last layer of the checkpoint dictionary. Otherwise you can pass in a list of layers to remove from the checkpoint before loading it (e.g. you would do that when loading an inception model that has more than one output layer).


pywick.models.model_utils.get_model(model_type: pywick.models.model_utils.ModelType, model_name: str, num_classes: int, pretrained: bool = True, force_reload: bool = False, custom_load_fn: Callable = None, **kwargs)[source]
  • model_type – (ModelType): type of model we’re trying to obtain (classification or segmentation)
  • model_name – (string): name of the model. By convention (for classification models) lowercase names represent pretrained model variants while Uppercase do not.
  • num_classes – (int): number of classes to initialize with (this will replace the last classification layer or set the number of segmented classes)
  • pretrained – (bool): whether to load the default pretrained version of the model NOTE! NOTE! For classification, the lowercase model names are the pretrained variants while the Uppercase model names are not. The only exception applies to torch.hub models (all efficientnet, mixnet, mobilenetv3, mnasnet, spnasnet variants) where a single lower-case string can be used for vanilla and pretrained versions. Otherwise, it is IN ERROR to specify an Uppercase model name variant with pretrained=True but one can specify a lowercase model variant with pretrained=False (default: True)
  • force_reload – (bool): Whether to force reloading the list of models from torch.hub. By default, a cache file is used if it is found locally and that can prevent new or updated models from being found.
  • custom_load_fn – (Callable): A custom callable function to use for loading models (typically used to load cutting-edge or custom models that are not in the publicly available list)


pywick.models.model_utils.get_fc_names(model_name, model_type=<ModelType.CLASSIFICATION: 'classification'>)[source]

Look up the name of the FC (fully connected) layer(s) of a model. Typically these are the layers that are replaced when transfer-learning from another model. Note that only a handful of models have more than one FC layer. Currently only ‘classification’ models are supported.

  • model_name – (string) name of the model
  • model_type – (ModelType) only classification is supported at this time

list names of the FC layers (usually a single one)

pywick.models.model_utils.get_supported_models(type: pywick.models.model_utils.ModelType)[source]
Parameters:type – (ModelType): classification or segmentation
Returns:list (strings) of supported models