查看.pth文件的网络结构与参数,输出模型参数个数
·
1、查看.pth文件的网络结构与参数
.pth文件结构
print("taking snapshot ...")
print("exp =", cfg.TRAIN.SNAPSHOT_DIR)
snapshot_dir = os.path.join(cfg.TRAIN.SNAPSHOT_DIR, f"{i_iter}.pth")
torch.save({'i_iter': i_iter,
'feature_extractor': feature_extractor.state_dict(),
'classifier':classifier.state_dict(),
'aux':aux.state_dict(),
'model_D': model_D.state_dict(),
'model_Dis': model_Dis.state_dict()
}, snapshot_dir)
查看网络层和参数
import torch
checkpoint = torch.load("/media/ailab/data/syn/Trans_depth2/depth_distribution/experiments/snapshots/SYNTHIA3Cityscapes_DeepLabv2_Depdis/75000.pth", map_location=torch.device('cpu'))
#torch.load('路径') 但是我的电脑没有GPU,是集成显卡呜呜呜,所以还得加个后面那部分map_location=torch.device('cpu')
print(checkpoint['classifier'].keys())
print(checkpoint['classifier'].shape)
2、输出模型参数个数
print("Total number of feature_extractor:{}".format(sum(x.numel() for x in feature_extractor.parameters())))
更多推荐


所有评论(0)