def calculate_models_params_and_flops():
backbone='darknet53'
network = backbones.get_backbone(backbone, pretrained=False, feature_type='tri_stage_fpn', n_classes=1000)
input = torch.randn(1, 3, 448, 448)
bb_macs, _ = profile(network, inputs=(input, ),verbose=False)
bb_params = summary(network,(3,448,448),device='cpu')[0].item()
model_df = pd.DataFrame(columns=['model_name', 'params', 'mac/flops'])
for key in model_argmap:
model_cfg= create_model_cfg(key, model_argmap)
model_components = create_model_components(model_cfg['name'],
preprocess_args=model_cfg['preprocess_args'],
network_args=model_cfg['network_args'],
loss_args=model_cfg['loss_args'],
postprocess_args=model_cfg['postprocess_args'],
stage='validate')
network=model_components.network
model_macs, model_params = profile(network, inputs=(input, ),verbose=False)
if 'backbone' in model_cfg.network_args:
model_macs-=bb_macs
model_params-=bb_params
key+='_head'
model_macs_in_g=model_macs/(10**9)
model_params_in_m=model_params/(10**6)
print(key,model_params_in_m,model_macs_in_g)
model_df = model_df.append({'model_name': key,
'params': model_params_in_m,
'mac/flops': model_macs_in_g}, ignore_index=True)
model_df.to_csv('models_summary.txt',index=False)
def test_create_model_components() :
preprocess_args, model_args = EasyDict(), EasyDict()
postprocess_args, loss_args = EasyDict(), EasyDict()
preprocess_args.input_size = 640
model_args.backbone = 'darknet53'
model_args.pyramid_channels = 256
loss_args.neg_pos = 3
loss_args.overlap_thresh = 0.5
model_name, stage = 'RetinaFace', 'train'
modules = models.create_model_components(
model_name=model_name,
preprocess_args=preprocess_args,
network_args=model_args,
loss_args=loss_args,
postprocess_args=postprocess_args,
stage=stage,
)
retinaface = modules.network
loss_fn = modules.loss
predictions = retinaface(torch.rand(1,3,640,640))
assert len(predictions) == 3
gt = [[
0.5, 0.5, 0.6, 0.6, 0.4, 0.4, 0.6, 0.6, 0.4, 0.4, 0.6, 0.6, 0.4, 0.4, 0
]]
gt = tensor(gt).unsqueeze(0)
loss = loss_fn(predictions, gt)
assert not any(torch.isinf(loss).view(-1)) and not any(torch.isnan(loss).view(-1))
retinaface.eval()
assert retinaface(torch.rand(1,3,640,640)).size(2) == 16
def test_darknet53_retinaface_predictor() :
preprocess_args, model_args = EasyDict(), EasyDict()
postprocess_args, loss_args = EasyDict(), EasyDict()
preprocess_args.input_size = 640
model_args.backbone = 'darknet53'
model_args.pyramid_channels = 256
loss_args.neg_pos = 3
loss_args.overlap_thresh = 0.5
model_name, stage = 'RetinaFace', 'train'
modules = models.create_model_components(
model_name=model_name,
preprocess_args=preprocess_args,
network_args=model_args,
loss_args=loss_args,
postprocess_args=postprocess_args,
stage=stage,
)
retinaface = modules.network
postprocess = modules.postprocess
predictor = BasePredictor(
model=retinaface,
postprocess=postprocess
).eval()
detections = predictor(
torch.rand(1,3,640,640),
tensor([0.01]),
tensor([0.01])
)
assert detections[0].dim() == 2
def export_model(model_name):
args = model_argmap[model_name]
model = create_model_components(model_name,
args['preprocess_args'],
args['network_args'],
loss_args=args['loss_args'],
postprocess_args=args['postprocess_args'])
predictor = create_predictor(model)
input_size = args['preprocess_args']['input_size']
output_path = os.path.join(output_dir, '{}.pt'.format(model_name))
exporter = TorchScriptExporter(output_path,
image_size=input_size,
check_tolerance=1e-6)
result = exporter(predictor)
assert os.path.exists(output_path)
return output_path
def test_model(task, backbone):
config_path = proj_path.joinpath("tests", "config",
"test_{}.yml".format(task))
config = load_config(config_path)
check_result = check_config(config, experiment_type='train')
assert check_result.valid, "config file %s for task %s is not valid, "\
"result:\n%s" % (config_path, task, str(check_result))
config.model.network_args.backbone = backbone
if backbones[0] == 'darknet53':
config.model.network_args.pretrained_backbone = None
args = {
'model_name': config.model.name,
'preprocess_args': config.model.preprocess_args,
'network_args': config.model.network_args,
'loss_args': config.model.loss_args,
'postprocess_args': config.model.postprocess_args,
'stage': 'train'
}
model = create_model_components(**args)
num_classes = config.model.network_args.n_classes
assert hasattr(model.network, "output_format"), "model {} doesn't have 'output_format' "\
"attribute explaining the output of the model".format(config.model.name)
x = torch.randn(1, 3, 640, 640)
x = model.network(x)
t = torch.tensor(0)
if task == 'classification':
t = torch.randint(0, num_classes, (1, ))
assert x.size() == torch.Size([1, num_classes]), \
"expected output size of %s for backbone '%s', got %s" % \
(torch.Size([1, num_classes]), backbones[0], x.size())
elif task == 'detection':
t = torch.tensor([[[14.0000, 0.4604, 0.0915, 0.2292, 0.3620],
[12.0000, 0.0896, 0.1165, 0.7583, 0.6617],
[14.0000, 0.1958, 0.2705, 0.0729, 0.0978]]])
assert len(
x) == 3, "expected output to have 3 elements, got %s" % len(x)
assert x[0].size(-1) == num_classes+5, "expected output model elements to have "\
"torch.Size([*, %s]), got %s" % (num_classes+5, x[0].size())
assert model.network.task == task
l = model.loss(x, t)
def test_exporter(model_name, image, backbone="resnet18", remove_output=True):
args = model_argmap[model_name]
args.network_args.backbone = backbone
model = create_model_components(model_name,
args['preprocess_args'],
args['network_args'],
loss_args=args['loss_args'],
postprocess_args=args['postprocess_args'])
predictor = create_predictor(model)
input_size = args['preprocess_args']['input_size']
output_path = os.path.join(output_dir,
'{}_{}.pt'.format(model_name, backbone))
output_format = predictor.output_format
additional_inputs = None
if hasattr(predictor.postprocess, 'additional_inputs'):
additional_inputs = predictor.postprocess.additional_inputs
print(" >> Exporting...")
exporter = TorchScriptExporter(output_path,
image_size=input_size,
check_tolerance=1e-6)
ok = exporter(predictor, example_image_path=image)
assert ok and os.path.exists(output_path)
del model, predictor, exporter
model = torch.jit.load(output_path)
for name, value in output_format.items():
assert hasattr(model, name + '_axis')
assert hasattr(model, name + '_indices')
assert getattr(model, name + '_axis') == torch.tensor(value['axis'])
assert all(
getattr(model, name +
'_indices') == torch.tensor(value['indices']))
inputs = (('input', (1, input_size, input_size, 3)), )
x = torch.randint(0, 256, (1, input_size, input_size, 3))
run_inputs, run_inputs_kwargs = [x], {}
if additional_inputs:
inputs += additional_inputs
for n, (name, shape) in enumerate(inputs):
assert hasattr(model, name + '_input_shape')
assert hasattr(model, name + '_input_pos')
assert getattr(model, name + '_input_shape').equal(torch.tensor(shape))
assert getattr(model, name + '_input_pos').equal(torch.tensor(n))
if name != 'input':
run_inputs.append(torch.zeros(*shape))
run_inputs_kwargs[name] = torch.zeros(*shape)
shape_name = [
name.replace('_input_shape', '')
for name, _ in model.named_buffers(recurse=False)
if name.endswith('_input_shape')
]
pos_name = [
name.replace('_input_pos', '')
for name, _ in model.named_buffers(recurse=False)
if name.endswith('_input_pos')
]
assert sorted(shape_name) == sorted(pos_name)
print(" >> Evaluating...")
with torch.no_grad():
out_eval = model(*run_inputs)
assert len(out_eval) == 1 ## single batch
assert out_eval[0].shape[-1] == sum(
len(v['indices'])
for v in output_format.values()) # number of elements
## TODO: check using keyword arguments, currently unsupported
# out_eval_kwargs = model(x, **run_inputs_kwargs)
# assert out_eval_kwargs.equal(out_eval)
if remove_output:
os.remove(output_path)
del model, out_eval
def create_model(model_config: EasyDict,
state_dict: Union[str, dict, Path] = None,
stage: str = 'train') -> EasyDict:
"""Function to create model and it's signature components. E.g. loss function, collate function, etc
Args:
model_config (EasyDict): Experiment file configuration at `model` section, as EasyDict object
state_dict (Union[str, dict, Path], optional): [description]. `model` Pytorch state dictionary or commonly known as weight, can be provided as the path to the file, or the returned dictionary object from `torch.load`. If this param is provided, it will override checkpoint specified in the experiment file. Defaults to None.
stage (str, optional): If set to 'train', this will enforce that the model must have `loss` and `collate_fn` attributes, hence it will make sure model can be used for training stage. If set to 'validate' it will ignore those requirements but cannot be used in training pipeline, but may still valid for other pipelines. Defaults to 'train'.
Raises:
TypeError: Raises if the provided `stage` not in 'train' or 'validate'
Returns:
EasyDict: The dictionary containing the model's components
Example:
The dictionary returned will contain several keys :
- `network` : Pytorch model's object which inherit `torch.nn.Module` class.
- `preprocess` : model's preprocessing module
- `postprocess` : model's postprocessing module
- `loss` : if provided, module for model's loss function
- `collate_fn` : if provided, module to be embedded to dataloader's `collate_fn` function to modify dataset label's format into desirable format that can be accepted by `loss` components
```python
from vortex.development.core.factory import create_model
from easydict import EasyDict
model_config = EasyDict({
'name': 'softmax',
'network_args': {
'backbone': 'efficientnet_b0',
'n_classes': 10,
'pretrained_backbone': True,
},
'preprocess_args': {
'input_size': 32,
'input_normalization': {
'mean': [0.4914, 0.4822, 0.4465],
'std': [0.2023, 0.1994, 0.2010],
'scaler': 255,
}
},
'loss_args': {
'reduction': 'mean'
}
})
model_components = create_model(
model_config = model_config
)
print(model_components.keys())
```
"""
if stage not in ['train', 'validate']:
raise TypeError(
'Unknown model "stage" argument, got {}, expected "train" or "validate"'
% stage)
logging.info('Creating Pytorch model from experiment file')
model_name = model_config.name
try:
preprocess_args = model_config.preprocess_args
except:
preprocess_args = {}
try:
network_args = model_config.network_args
except:
network_args = {}
try:
loss_args = model_config.loss_args
except:
loss_args = {}
try:
postprocess_args = model_config.postprocess_args
except:
postprocess_args = {}
model_components = create_model_components(
model_name,
preprocess_args=preprocess_args,
network_args=network_args,
loss_args=loss_args,
postprocess_args=postprocess_args,
stage=stage)
if not isinstance(model_components, EasyDict):
model_components = EasyDict(model_components)
if 'init_state_dict' in model_config or state_dict is not None:
if isinstance(state_dict, Path):
state_dict = str(state_dict)
model_path = None
# Load state_dict from config if specified in experiment file
if 'init_state_dict' in model_config and state_dict is None:
logging.info(
"Loading state_dict from configuration file : {}".format(
model_config.init_state_dict))
model_path = model_config.init_state_dict
# If specified using function's parameter, override the experiment config init_state_dict
elif isinstance(state_dict, str):
logging.info("Loading state_dict : {}".format(state_dict))
model_path = state_dict
if ('init_state_dict' in model_config
and state_dict is None) or isinstance(state_dict, str):
assert model_path
ckpt = torch.load(model_path)
state_dict = ckpt['state_dict'] if 'state_dict' in ckpt else ckpt
assert isinstance(state_dict, (OrderedDict, dict))
model_components.network.load_state_dict(state_dict, strict=True)
return model_components