fix keyerror?
This commit is contained in:
parent
0abd4b94fb
commit
e838f591aa
|
@ -5,9 +5,11 @@ def estimate_model_size(config: dict):
|
||||||
:param config:
|
:param config:
|
||||||
:return:
|
:return:
|
||||||
"""
|
"""
|
||||||
vocab_size = config['vocab_size']
|
vocab_size = config.get('vocab_size')
|
||||||
hidden_size = config['hidden_size']
|
hidden_size = config.get('hidden_size')
|
||||||
num_hidden_layers = config['num_hidden_layers']
|
num_hidden_layers = config.get('num_hidden_layers')
|
||||||
intermediate_size = config['intermediate_size']
|
intermediate_size = config.get('intermediate_size')
|
||||||
|
if vocab_size and hidden_size and num_hidden_layers and intermediate_size:
|
||||||
total_params = (vocab_size * hidden_size) + (num_hidden_layers * ((hidden_size * intermediate_size * 4) + (hidden_size * hidden_size * 3)))
|
total_params = (vocab_size * hidden_size) + (num_hidden_layers * ((hidden_size * intermediate_size * 4) + (hidden_size * hidden_size * 3)))
|
||||||
return int(total_params / 1e9)
|
return int(total_params / 1e9)
|
||||||
|
return 0
|
||||||
|
|
Reference in New Issue