pytorch常用功能代码汇总

统计网络参数量

# 所有参数量
p_all = sum([torch.numel(param) for param in model.parameters()])
# 需进行梯度更新的参数量
p_require_grad = sum([torch.numel(param) for param in model.parameters() if param.requires_grad])

统计optimizer中参数量

sum([torch.numel(param) for param in optimizer.param_groups[0]['params'] if param.requires_grad])	# 这里只考虑只有一组param

冻结网络部分参数

冻结：

for param in model.parameters():	# 设置对应的参数requires_grad = False
	param.requires_grad = False
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr, [beta1, beta2])	# 在optimizer中传入需要更新梯度的参数

解冻：

for param in model.parameters():	# 设置对应的参数requires_grad = True
	param.requires_grad = True
optimizer = optim.Adam(model.parameters(), lr, [beta1, beta2])	# 在optimizer中传入模型的参数