def _vgg(cfg, pretrained=False, progress=True, num_classes=100, device='cuda'): model = VGG(cfg, num_classes=num_classes) if pretrained: checkpoint_url = model_urls[cfg] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _mobilenetv2_imagenet10(arch, alpha=1.0, pretrained=False, progress=True, num_classes=10, device='cuda'): model = torchvision.models.mobilenet.MobileNetV2(width_mult=alpha, num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def unet_carvana(pretrained=False, progress=True, num_classes=1, device="cuda"): model = UNet(n_channels=3, n_classes=num_classes, bilinear=True) if pretrained: checkpoint_url = model_urls["unet_carvana"] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def resnet50_tinyimagenet(pretrained=False, progress=True, num_classes=100, device="cuda"): model = models.resnet50(num_classes=num_classes) if pretrained: checkpoint_url = model_urls['resnet50'] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def yolo3( net="yolo3", dataset_name="voc", num_classes=20, pretrained=False, progress=True, device="cuda", **kwargs ): config_path = get_project_root() / CFG_PATH / yolov3_cfg[net] model = YoloV5_6(config_path, ch=3, nc=num_classes) if pretrained: checkpoint_url = urlparse.urljoin(CHECKPOINT_STORAGE_URL, model_urls[f"{net}_{dataset_name}"]) model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _mobilenetv1_vww(arch, pretrained=False, progress=True, num_classes=2, device='cuda'): model = MobileNetV1(num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _shufflenetv2(arch, net_size=1, pretrained=False, num_classes=100, progress=True, device='cuda'): model = ShuffleNetV2(net_size, num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _resnet_imagenet10(arch, pretrained=False, progress=True, num_classes=10, device='cuda'): model = torchvision.models.resnet18(num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _lenet_mnist(arch, pretrained=False, progress=True, num_classes=10, device="cuda"): model = LeNet5(output=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _pre_act_resnet(arch, block, layers, num_classes=100, pretrained=False, progress=True, device='cuda'): model = PreActResNet(block, layers, num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _mlp10_mnist(arch, n_hiddens, pretrained=False, progress=True, num_classes=10, device="cuda"): model = MLP(input_dims=784, n_hiddens=n_hiddens, n_class=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _densenet(arch, block, layers, growth_rate=32, pretrained=False, num_classes=100, progress=True, device='cuda'): model = DenseNet(block, layers, growth_rate, num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def fcn32( net="fcn32", dataset="voc", num_classes=21, pretrained=False, progress=True, device="cuda", ): model = FCN(n_class=num_classes) if pretrained: checkpoint_url = model_urls[f"{net}_{dataset}"] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def deeplab( backbone="resnet", dataset="voc", num_classes=21, pretrained=False, progress=True, device="cuda", ): model = DeepLab(backbone=backbone, num_classes=num_classes) if pretrained: checkpoint_url = model_urls[f"{backbone}_{dataset}"] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _mobilenetv3_vww(arch="small", pretrained=False, progress=True, num_classes=2, device='cuda'): if arch == "small": model = mobilenetv3_small(num_classes=num_classes) elif arch == "large": model = mobilenetv3_large(num_classes=num_classes) if pretrained: checkpoint_url = model_urls[f"mobilenetv3_{arch}"] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _resnet_imagenet16(arch, pretrained=False, progress=True, num_classes=16, device='cuda'): if arch == "resnet18": model = torchvision.models.resnet18(num_classes=num_classes) elif arch == "resnet50": model = torchvision.models.resnet50(num_classes=num_classes) else: raise ValueError if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def ssd_model(model_name, dataset, create_model_fn, config, num_classes=20, pretrained=False, progress=True, device='cuda'): model = create_model_fn(num_classes + 1) # + 1 for background class model.config = config model.priors = config.priors.to(device) if pretrained: checkpoint_url = urlparse.urljoin( CHECKPOINT_STORAGE_URL, model_urls[f"{model_name}_{dataset}"]) model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _resnext(arch, num_blocks, cardinality, bottleneck_width, num_classes=100, pretrained=False, progress=True, device='cuda'): model = ResNeXt( num_blocks=num_blocks, cardinality=cardinality, bottleneck_width=bottleneck_width, num_classes=num_classes, ) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def yolo5_6(net="yolo5_6s", dataset_name="voc", num_classes=20, activation_type=None, pretrained=False, progress=True, device="cuda"): config_key = net for suffix in MODEL_NAME_SUFFICES: config_key = re.sub(f'\_{suffix}$', '', config_key) # pylint: disable=W1401 config_path = get_project_root() / CFG_PATH / yolov5_cfg[config_key] model = YoloV5_6(config_path, ch=3, nc=num_classes, activation_type=activation_type) if pretrained: checkpoint_url = urlparse.urljoin(CHECKPOINT_STORAGE_URL, model_urls[f"{net}_{dataset_name}"]) model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def unet_enc_dec( enc_type="resnet50", dec_type="unet_scse", output_channels=21, dataset_type="voc", num_filters=16, pretrained=True, progress=False, device="cuda", ): model = EncoderDecoderNet( output_channels=output_channels, enc_type=enc_type, dec_type=dec_type, num_filters=num_filters, ) if pretrained: checkpoint_url = model_urls[f"{dec_type}_{enc_type}_{dataset_type}"] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)
def _googlenet(arch, pretrained=False, progress=True, num_classes=100, device='cuda'): model = GoogLeNet(num_classes=num_classes) if pretrained: checkpoint_url = model_urls[arch] model = load_pretrained_weights(model, checkpoint_url, progress, device) return model.to(device)