def test_ssd_resnsets(
key: str,
pretrained: Union[bool, str],
pretrained_backbone: Union[bool, str],
test_input: bool,
match_const: Callable,
):
model = ModelRegistry.create(key, pretrained)
diff_model = match_const(pretrained_backbone=pretrained_backbone)
if pretrained:
compare_model(model, diff_model, same=False)
match_model = ModelRegistry.create(key, pretrained)
compare_model(model, match_model, same=True)
if pretrained_backbone and pretrained_backbone is not True:
compare_model(model.feature_extractor, diff_model.feature_extractor, same=False)
match_model = ModelRegistry.create(key, pretrained_backbone=pretrained_backbone)
compare_model(
diff_model.feature_extractor, match_model.feature_extractor, same=True
)
if test_input:
input_shape = ModelRegistry.input_shape(key)
batch = torch.randn(1, *input_shape)
model.eval()
boxes, scores = model(batch)
assert isinstance(boxes, torch.Tensor)
assert isinstance(scores, torch.Tensor)
assert boxes.dim() == 3
assert scores.dim() == 3
assert boxes.size(0) == 1
assert boxes.size(1) == 4
assert scores.size(0) == 1
assert boxes.size(2) == scores.size(2) # check same num default boxes
def test_yolo_v3(
key: str,
pretrained: Union[bool, str],
pretrained_backbone: Union[bool, str],
test_input: bool,
match_const: Callable,
):
model = ModelRegistry.create(key, pretrained)
diff_model = match_const(pretrained_backbone=pretrained_backbone)
if pretrained:
compare_model(model, diff_model, same=False)
match_model = ModelRegistry.create(key, pretrained)
compare_model(model, match_model, same=True)
if pretrained_backbone and pretrained_backbone is not True:
compare_model(model.backbone, diff_model.backbone, same=False)
match_model = ModelRegistry.create(
key, pretrained_backbone=pretrained_backbone)
compare_model(diff_model.backbone, match_model.backbone, same=True)
if test_input:
input_shape = ModelRegistry.input_shape(key)
batch = torch.randn(1, *input_shape)
model.eval()
outputs = model(batch)
assert isinstance(outputs, list)
for output in outputs:
assert isinstance(output, torch.Tensor)
assert output.dim() == 5
assert output.size(-1) == 85
def create_model(args: Any, num_classes: int) -> Module:
"""
:param args: object with configuration for model classes
:param num_classes: Integer representing the number of output classes
:returns: A Module object representing the created model
"""
with torch_distributed_zero_first(
args.local_rank): # only download once locally
if args.checkpoint_path == "zoo":
if args.recipe_path and args.recipe_path.startswith("zoo:"):
args.checkpoint_path = Zoo.download_recipe_base_framework_files(
args.recipe_path, extensions=[".pth"])[0]
else:
raise ValueError(
"'zoo' provided as --checkpoint-path but a SparseZoo stub"
" prefixed by 'zoo:' not provided as --recipe-path")
model = ModelRegistry.create(
args.arch_key,
args.pretrained,
args.checkpoint_path,
args.pretrained_dataset,
num_classes=num_classes,
**args.model_kwargs,
)
print(f"created model: {model}")
return model
def test_mnist(key: str, pretrained: Union[bool, str], test_input: bool):
model = ModelRegistry.create(key, pretrained)
diff_model = mnist_net()
if pretrained:
compare_model(model, diff_model, same=False)
match_model = ModelRegistry.create(key, pretrained)
compare_model(model, match_model, same=True)
if test_input:
input_shape = ModelRegistry.input_shape(key)
batch = torch.randn(1, *input_shape)
out = model(batch)
assert isinstance(out, tuple)
for tens in out:
assert tens.shape[0] == 1
assert tens.shape[1] == 10
def test_torchvision_registry_models(key: str, pretrained: Union[bool, str],
constructor: Callable):
model = ModelRegistry.create(key, pretrained)
diff_model = constructor(pretrained=False)
compare_model(model, diff_model, same=False)
if pretrained is True:
# check torchvision weights are properly loaded
match_model = constructor(pretrained=pretrained)
compare_model(model, match_model, same=True)
def test_efficientnet(
key: str,
pretrained: Union[bool, str],
test_input: bool,
match_const: Callable,
model_args: dict,
):
model = ModelRegistry.create(key, pretrained, **model_args)
diff_model = match_const(**model_args)
if pretrained:
compare_model(model, diff_model, same=False)
match_model = ModelRegistry.create(key, pretrained, **model_args)
compare_model(model, match_model, same=True)
if test_input:
input_shape = ModelRegistry.input_shape(key)
batch = torch.randn(1, *input_shape)
model = model.eval()
out = model(batch)
assert isinstance(out, tuple)
for tens in out:
assert tens.shape[0] == 1
assert tens.shape[1] == 1000