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