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)
