关于thop踩坑

在使用老的 ultralytics yolov3包的时候,踩了thop的一个坑:

#具体位置在model.py
class Darknet(nn.Module):
    # YOLOv3 object detection model

    def __init__(self, cfg, img_size=(416, 416), verbose=False):
        super(Darknet, self).__init__()

        self.module_defs = parse_model_cfg(cfg)
        self.module_list, self.routs = create_modules(self.module_defs, img_size, cfg)  # self.routs是以二进制记录的是
        # 否有跳层连接的list,其中对应位置提供跳层连接的图像数据的层
        self.yolo_layers = get_yolo_layers(self) # 获取yolo层的位置,返回yolo层在model里的的index
        # torch_utils.initialize_weights(self)

        # Darknet Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
        self.version = np.array([0, 2, 5], dtype=np.int32)  # (int32) version info: major, minor, revision
        self.seen = np.array([0], dtype=np.int64)  # (int64) number of images seen during training
		# 下面这一句会调用thop的profile函数
        self.info(verbose) if not ONNX_EXPORT else None  # print model description
...
    def info(self, verbose=False):
        torch_utils.model_info(self, verbose)

# 从torch_utils里把这段复制过来了
def model_info(model, verbose=False):
...

    try:  # FLOPS
        from thop import profile
        # 重点在这 
        macs, _ = profile(model, inputs=(torch.zeros(1, 3, 416, 416),), verbose=True)

# 而在profile里会运行一次forward
    with torch.no_grad():
        model(*inputs)

在源代码中model在创建的时候就进行了一次forward,而此时model仍在cpu上,这就造成了一系列tensor创建时在cpu上,而后在训练测试时没有随着parameter转移到gpu上的问题

解决方案

直接注释掉init里的这一行,需要用的时候再自行info
(这一行代码也有可能是我自己加的,由于没用git做代码版本管理,有些混乱)

你可能感兴趣的:(python,yolov3,深度学习,python,pytorch)