对比SPP和SPPF的计算结果以及速度

借鉴的大佬的代码和讲解
https://blog.csdn.net/qq_37541097/article/details/123594351

代码:

#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@File  : SPP_SPPF_compare
@Author: 
@Time  : 2023/1/11 9:54
@Desc  : 对比下SPP和SPPF的计算结果以及速度
这里将SPPF中最开始和结尾处的1x1卷积层给去掉了,只对比含有MaxPool的部分
"""

import time
import torch
import torch.nn as nn


class SPP(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool1 = nn.MaxPool2d(5, 1, padding=2)
        self.maxpool2 = nn.MaxPool2d(9, 1, padding=4)
        self.maxpool3 = nn.MaxPool2d(13, 1, padding=6)

    def forward(self, x):
        o1 = self.maxpool1(x)
        o2 = self.maxpool2(x)
        o3 = self.maxpool3(x)
        return torch.cat([x, o1, o2, o3], dim=1)


class SPPF(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool = nn.MaxPool2d(5, 1, padding=2)

    def forward(self, x):
        o1 = self.maxpool(x)
        o2 = self.maxpool(o1)
        o3 = self.maxpool(o2)
        return torch.cat([x, o1, o2, o3], dim=1)


def main():
    input_tensor = torch.rand(8, 32, 16, 16)
    spp = SPP()
    sppf = SPPF()
    output1 = spp(input_tensor)
    output2 = sppf(input_tensor)

    print(torch.equal(output1, output2))  # 比较两个张量的形状和各个元素是否都相等.

    t_start = time.time()
    for _ in range(100):
        spp(input_tensor)
    print(f"spp time: {time.time() - t_start}")

    t_start = time.time()
    for _ in range(100):
        sppf(input_tensor)
    print(f"sppf time: {time.time() - t_start}")


if __name__ == '__main__':
    main()

结果:

D:\anaconda\python.exe D:/98project/yolov5-7.0/tool/SPP_SPPF_compare.py
True
spp time: 0.8605170249938965
sppf time: 0.3405332565307617

Process finished with exit code 0

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