pytorch, tensorflow, keras统计模型参数大小

统计模型大小的方法统一步骤

• 1 统计总的参数个数，
• 2 利用参数个数算出权重大小

pytorch

def get_model_size(model):
	para_num = sum([p.numel() for p in model.parameters()])
	# para_size: 参数个数 * 每个4字节(float32) / 1024 / 1024，单位为 MB
	para_size = para_num * 4 / 1024 / 1024
	return para_size

tensorflow

import numpy as np
import tensorflow as tf

def get_model_size(model):
	para_num = sum([np.prod(var.get_shape().as_list()) for var in tf.trainable_variables()])
	# para_size: 参数个数 * 每个4字节(float32) / 1024 / 1024，单位为 MB
	para_size = para_num * 4 / 1024 / 1024
	return para_size

keras

import numpy as np

def get_model_size(model):
	para_num = sum([np.prod(w.shape) for w in model.get_weights()])
	# para_size: 参数个数 * 每个4字节(float32) / 1024 / 1024，单位为 MB
	para_size = para_num * 4 / 1024 / 1024
	return para_size