查看.pth文件的网络结构与参数,输出模型参数个数_查看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())))