v1.5

README.md

@@ -52,7 +52,7 @@ image classifier implement in pytoch.
 - **简单的安装过程**  
   1. 安装好pytorch, torchvision(pytorch==1.12.0+ torchvision==0.13.0+)  
     可以在[pytorch](https://pytorch.org/get-started/previous-versions/)官网找到对应的命令进行安装.
-  2. pip install -r requirements.txt
+  2. pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
 
 - **人性化的设定**
   1. 大部分可视化数据(混淆矩阵,tsne,每个类别的指标)都会以csv或者log的格式保存到本地,方便后期美工图像.
@@ -245,7 +245,7 @@ image classifier implement in pytoch.
     是否采用测试阶段的数据增强.  
   - **cam_visual**  
     default: False  
-    是否进行热力图可视化.  
+    是否进行热力图可视化.(目前只支持在cpu上进行可视化热力图)  
   - **cam_type**  
     type: string, default: GradCAMPlusPlus, choices: ['GradCAM', 'HiResCAM', 'ScoreCAM', 'GradCAMPlusPlus', 'AblationCAM', 'XGradCAM', 'EigenCAM', 'FullGrad']  
     热力图可视化的类型.  
@@ -533,7 +533,7 @@ image classifier implement in pytoch.
 
     1. 在predict.py中的--cam_type参数,可以选择各种各样的热力图计算方法,如果遇到这个效果不好可以进行更换,当然效果不好除了跟热力图的计算方法有关系之外还有是你选择的层.
     2. 在第一点中说到的层,热力图可视化需要选择层进行计算,这个是在predict.py中的cam_visual方法定义中进行定义.
-      例如这样 cam_model = cam_visual(model, test_transform, DEVICE, model.cam_layer(), opt) 这个model.cam_layer()就是选择的层,默认就是卷积特征的最后一个block,这个可以在model文件夹中的每个模型文件中的模型进行修改.
+      例如这样 cam_model = cam_visual(model, test_transform, DEVICE, model.cam_layer(), opt) 这个model.cam_layer()就是选择的层,默认就是卷积特征的最后一个layer,这个可以在model文件夹中的每个模型文件中的模型进行修改.
 
   13. 关于类别平衡的问题.
 
@@ -679,4 +679,11 @@ image classifier implement in pytoch.
 
 1. predict.py支持检测灰度图,其读取后会检测是否为RGB通道,不是的话会进行转换.
 2. 更新readme.md.
-3. 修复一些bug.
+3. 修复一些bug.
+
+### pytorch-classifier v1.5 更新日志
+
+1. 更新readme.md
+2. 修改predict.py中的读取模型方式.
+3. 修复predict.py中批量预测保存图片时候的内存溢出问题.
+4. 修复predict.py中热力图可视化的报错问题,但是目前只支持cpu状态下可视化热力图.

model/efficientnetv2.py

@@ -383,7 +383,7 @@ def forward_features(self, x, need_fea=False):
             return x
 
     def cam_layer(self):
-        return self.features[-1]
+        return self.features[-5:]
 
     def switch_to_deploy(self):
         new_block = []

predict.py

@@ -6,8 +6,10 @@
 os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
 import matplotlib.pyplot as plt
 import numpy as np
+from copy import deepcopy
 from utils import utils_aug
 from utils.utils import predict_single_image, cam_visual, dict_to_PrettyTable, select_device, model_fuse
+from utils.utils_model import select_model
 
 def set_seed(seed):
     random.seed(seed)
@@ -27,15 +29,21 @@ def parse_opt():
     parser.add_argument('--device', type=str, default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
 
     opt = parser.parse_known_args()[0]
+
     if not os.path.exists(os.path.join(opt.save_path, 'best.pt')):
         raise Exception('best.pt not found. please check your --save_path folder')
     ckpt = torch.load(os.path.join(opt.save_path, 'best.pt'))
     DEVICE = select_device(opt.device)
     if opt.half and DEVICE.type == 'cpu':
         raise Exception('half inference only supported GPU.')
     if opt.half and opt.cam_visual:
+        raise Exception('cam visual only supported cpu. please set device=cpu.')
+    if (opt.device != 'cpu') and opt.cam_visual:
         raise Exception('cam visual only supported FP32.')
-    model = ckpt['model'].float()
+    with open(opt.label_path) as f:
+        CLASS_NUM = len(f.readlines())
+    model = select_model(ckpt['model'].name, CLASS_NUM)
+    model.load_state_dict(ckpt['model'].float().state_dict(), strict=False)
     model_fuse(model)
     model = (model.half() if opt.half else model)
     model.to(DEVICE)
@@ -59,7 +67,7 @@ def parse_opt():
     opt, DEVICE, model, test_transform, label = parse_opt()
 
     if opt.cam_visual:
-        cam_model = cam_visual(model, test_transform, DEVICE, model.cam_layer(), opt)
+        cam_model = cam_visual(model, test_transform, DEVICE, opt)
 
     if os.path.isdir(opt.source):
         save_path = os.path.join(opt.save_path, 'predict', datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S'))
@@ -71,14 +79,16 @@ def parse_opt():
 
             plt.figure(figsize=(6, 6))
             if opt.cam_visual:
-                cam_output = cam_model(os.path.join(opt.source, file), pred)
+                cam_output = cam_model(os.path.join(opt.source, file))
                 plt.imshow(cam_output)
             else:
                 plt.imshow(plt.imread(os.path.join(opt.source, file)))
             plt.axis('off')
             plt.title('predict label:{}\npredict probability:{:.4f}'.format(label[pred], float(pred_result[pred])))
             plt.tight_layout()
             plt.savefig(os.path.join(save_path, file))
+            plt.clf()
+            plt.close()
 
         with open(os.path.join(save_path, 'result.csv'), 'w+') as f:
             f.write('img_path,pred_class,pred_class_probability\n')

utils/utils.py

@@ -16,6 +16,7 @@
 from sklearn.manifold import TSNE
 from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients
 from pytorch_grad_cam.utils.image import show_cam_on_image
 from collections import OrderedDict
 from .utils_aug import rand_bbox
@@ -655,22 +656,22 @@ def predict_single_image(path, model, test_transform, DEVICE, half=False):
     return int(pred_result.argmax()), pred_result
 
 class cam_visual:
-    def __init__(self, model, test_transform, DEVICE, target_layers, opt):
+    def __init__(self, model, test_transform, DEVICE, opt):
         self.test_transform = test_transform
         self.DEVICE = DEVICE
         self.opt = opt
 
-        self.cam_model = eval(opt.cam_type)(model=deepcopy(model), target_layers=[target_layers], use_cuda=torch.cuda.is_available())
+        self.cam_model = eval(opt.cam_type)(model=model, target_layers=model.cam_layer(), use_cuda=(DEVICE.type != 'cpu'))
+        self.model = model
 
-    def __call__(self, path, label):
+    def __call__(self, path):
         pil_img = Image.open(path)
         tensor_img = self.test_transform(pil_img).unsqueeze(0).to(self.DEVICE)
-        targets = [ClassifierOutputTarget(label)]
-
+
         if len(tensor_img.shape) == 5:
-            grayscale_cam_list = [self.cam_model(input_tensor=tensor_img[:, i], targets=targets) for i in range(tensor_img.size(1))]
+            grayscale_cam_list = [self.cam_model(input_tensor=tensor_img[:, i], targets=None) for i in range(tensor_img.size(1))]
         else:
-            grayscale_cam_list = [self.cam_model(input_tensor=tensor_img, targets=targets)]
+            grayscale_cam_list = [self.cam_model(input_tensor=tensor_img, targets=None)]
 
         grayscale_cam = np.concatenate(grayscale_cam_list, 0).mean(0)
         grayscale_cam = cv2.resize(grayscale_cam, pil_img.size)

utils/utils_model.py

@@ -3,7 +3,7 @@
 import model as models
 from thop import clever_format, profile
 
-def select_model(name, num_classes, input_shape, channels, pretrained=False):
+def select_model(name, num_classes, input_shape=None, channels=None, pretrained=False):
     if 'shufflenet_v2' in name:
         model = eval('models.{}(pretrained={})'.format(name, pretrained))
         model.fc = nn.Sequential(
@@ -118,22 +118,22 @@ def select_model(name, num_classes, input_shape, channels, pretrained=False):
     else:
         raise 'Unsupported Model Name.'
 
+    if input_shape and channels:
     # 计算参数量和flops
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    dummy_input = torch.randn(1, channels, input_shape[0], input_shape[1]).to(device)
-    flops, params = profile(model.to(device), (dummy_input,), verbose=False)
-    #--------------------------------------------------------#
-    #   flops * 2是因为profile没有将卷积作为两个operations
-    #   有些论文将卷积算乘法、加法两个operations。此时乘2
-    #   有些论文只考虑乘法的运算次数，忽略加法。此时不乘2
-    # --------------------------------------------------------#
-    # flops = flops * 2
-    flops, params = clever_format([flops, params], "%.3f")
-    print('Select Model: {}'.format(name))
-    print('Total FLOPS: %s' % (flops))
-    print('Total params: %s' % (params))
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        dummy_input = torch.randn(1, channels, input_shape[0], input_shape[1]).to(device)
+        flops, params = profile(model.to(device), (dummy_input,), verbose=False)
+        #--------------------------------------------------------#
+        #   flops * 2是因为profile没有将卷积作为两个operations
+        #   有些论文将卷积算乘法、加法两个operations。此时乘2
+        #   有些论文只考虑乘法的运算次数，忽略加法。此时不乘2
+        # --------------------------------------------------------#
+        # flops = flops * 2
+        flops, params = clever_format([flops, params], "%.3f")
+        print('Select Model: {}'.format(name))
+        print('Total FLOPS: %s' % (flops))
+        print('Total params: %s' % (params))
     model.name = name
-
     return model
 
 if __name__ == '__main__':