pytorch从零开始YOLOV4的neck部分代码写在yolo4.py里
整体代码如下:
import torch # 自己写的代码
import torch.nn as nn
from collections import OrderedDict
from torch.autograd import Variable
from CSPDarknet import *
import cv2
import math
# CBL的构建
def conv2d(filter_in, filter_out, kernel_size, stride=1):
pad = (kernel_size - 1) // 2 if kernel_size else 0 # 如果kernel_size为1,pad就为0,如果kernel_size为3,pad为1
return nn.Sequential(OrderedDict([
("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=stride, padding=pad, bias=False)),
("bn", nn.BatchNorm2d(filter_out)),
("relu", nn.LeakyReLU(0.1)),
]))
# ------------------------------------------------#
# SPP结构,利用不同大小的池化核继续宁池化
# 池化后堆叠
# ------------------------------------------------#
class SpatialPyramidPooling(nn.Module):
def __init__(self, pool_sizes=[5, 9, 13]):
super(SpatialPyramidPooling, self).__init__()
self.maxpools = nn.ModuleList([nn.MaxPool2d(pool_size, 1, pool_size // 2) for pool_size in pool_sizes])
def forward(self, x):
features = [maxpool(x) for maxpool in self.maxpools[