想起来就更新....
yolov5-6.1/models/yolo.py
class Detect(nn.Module):
stride = None # strides computed during build
onnx_dynamic = False # ONNX export parameter
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
super().__init__()
self.nc = nc # number of classes
self.no = nc + 5 # number of outputs per anchor
self.nl = len(anchors) # number of detection layers
self.na = len(anchors[0]) // 2 # number of anchors
self.grid = [torch.zeros(1)] * self.nl # init grid
self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
self.inplace = inplace # use in-place ops (e.g. slice assignment)
小提示:
模型中需要保存下来的参数包括两种:一种是反向传播需要被 optimizer 更新的,称为 parameter;一种是反向传播不要被 optimizer 更新,称为 buffer。第二种参数需要创建 tensor,然后将 tensor 通过 register_buffer() 进行注册,可以通过 model.buffers() 返回,注册完后参数也会自动保存到 orderdict 有序字典中去。注意:buffer 的更新在 forward 中,optim.step 只能更新 nn.parameter 类型的参数。