本文重點描述開源YoloV5代碼實現的細節，不會對YoloV5的整體思路進行介紹，整體思路可以參考江大白的博客

江大白：深入淺出Yolo系列之Yolov3&Yolov4&Yolov5&Yolox核心基礎知識完整講解2644 贊同 · 332 評論文章正在上傳…重新上傳取消

講解的很細致，建議閱讀之后再來看本篇文章。

聲明：本文有些圖摘自江大白的上述博客，如有侵權，請聯系本人刪除

本文所使用的代碼為2021-08-23日flok的官網YoloV5代碼倉

GitHub - xuanzhangyang/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite?github.com/xuanzhangyang/yolov5正在上傳…重新上傳取消

(又發現一篇還不錯的介紹?進擊的后浪yolov5深度可視化解析 - 知乎)

一、數據集相關代碼解讀

create_dataloader函數代碼解讀(utils/datasets.py)

該函數在train.py的205和215行調用，分別用來創建訓練數據集的loader和評估數據集的loader。

def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0,rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix=''):# Make sure only the first process in DDP process the dataset first, and the following others can use the cachewith torch_distributed_zero_first(rank):

torch_distributed_zero_first函數的作用是只有主進程來加載數據，其他進程處于等待狀態直到主進程加載完數據，該函數具體實現說明參考下面的torch_distributed_zero_first函數解讀

dataset = LoadImagesAndLabels(path, imgsz, batch_size,augment=augment, # augment imageshyp=hyp, # augmentation hyperparametersrect=rect, # rectangular trainingcache_images=cache,single_cls=single_cls,stride=int(stride),pad=pad,image_weights=image_weights,prefix=prefix)

LoadImagesAndLabels函數加載數據集，該函數具體實現說明參考下面的LoadImagesAndLabels類代碼解讀

batch_size = min(batch_size, len(dataset))nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, workers]) # number of workerssampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else Noneloader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader# Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()dataloader = loader(dataset,batch_size=batch_size,num_workers=nw,sampler=sampler,pin_memory=True,collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)return dataloader, dataset

torch_distributed_zero_first函數解讀(utils/torch_utils.py)

pytorch在分布式訓練過程中，對于數據的讀取是采用主進程預讀取并緩存，然后其它進程從緩存中讀取，不同進程之間的數據同步具體通過torch.distributed.barrier()實現。

def torch_distributed_zero_first(local_rank: int):"""Decorator to make all processes in distributed training wait for each local_master to do something."""if local_rank not in [-1, 0]:dist.barrier(device_ids=[local_rank])yieldif local_rank == 0:dist.barrier(device_ids=[0])

torch_distributed_zero_first是在create_dataloader函數中調用的，如果執行create_dataloader()函數的進程不是主進程，即rank不等于0或者-1，上下文管理器會執行相應的torch.distributed.barrier()，設置一個阻塞柵欄，讓此進程處于等待狀態，等待所有進程到達柵欄處（包括主進程數據處理完畢）；

如果執行create_dataloader()函數的進程是主進程，其會直接去讀取數據并處理，然后其處理結束之后會接著遇到torch.distributed.barrier()，此時，所有進程都到達了當前的柵欄處，這樣所有進程就達到了同步，并同時得到釋放。

LoadImagesAndLabels類代碼解讀(utils/datasets.py)

該類繼承pytorch的Dataset類，需要實現父類的__init__方法， __getitem__方法和__len__方法， 在每個step訓練的時候，DataLodar迭代器通過__getitem__方法獲取一批訓練數據。

__init__函數解讀

def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):...self.albumentations = Albumentations() if augment else Nonetry:f = [] # image filesfor p in path if isinstance(path, list) else [path]:p = Path(p) # os-agnosticif p.is_dir(): # dirf += glob.glob(str(p / '**' / '*.*'), recursive=True)# f = list(p.rglob('**/*.*')) # pathlibelif p.is_file(): # filewith open(p, 'r') as t:t = t.read().strip().splitlines()parent = str(p.parent) + os.sepf += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path# f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)else:raise Exception(f'{prefix}{p} does not exist')self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS])

self.img_files里面存放的就是所有的圖片的路徑， 并且是排序好的， 類似[‘coco128/images/train2017/000000000009.jpg’, ‘coco128/images/train2017/000000000025.jpg’, ‘coco128/images/train2017/000000000030.jpg’]

# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlibassert self.img_files, f'{prefix}No images found'except Exception as e:raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')# Check cacheself.label_files = img2label_paths(self.img_files) # labels

通過圖像的路徑找到圖片對應的標注文件的路徑，類似[‘coco128/labels/train2017/000000000009.txt’, ‘coco128/labels/train2017/000000000025.txt’, ‘coco128/labels/train2017/000000000030.txt’]

cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')

catch_path： ‘coco128/labels/train2017.cache’

try:cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dictassert cache['version'] == 0.4 and cache['hash'] == get_hash(self.label_files + self.img_files)except:cache, exists = self.cache_labels(cache_path, prefix), False # cache

cache是個字典，key的個數為圖像個數+4，存儲了每張圖像對應的所有gt box標簽和圖像寬高，以及cache的hash值等不太重要的信息， 存儲的標簽如下：

(Pdb) p cache['coco128/images/train2017/000000000625.jpg'] [array([[ 0, 0.725, 0.69758, 0.1875, 0.47982],[ 0, 0.50391, 0.64675, 0.15781, 0.60762],[ 0, 0.36186, 0.73258, 0.14428, 0.41798],[ 29, 0.45219, 0.27519, 0.052734, 0.048498]], dtype=float32), (640, 446), []]

.

# Display cachenf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, totalif exists:d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"tqdm(None, desc=prefix + d, total=n, initial=n) # display cache resultsif cache['msgs']:logging.info('\n'.join(cache['msgs'])) # display warningsassert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}'# Read cache[cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove itemslabels, shapes, self.segments = zip(*cache.values())

label就是gt box的信息，包括類別的坐標。 shapes是圖像寬高信息。 segments都是空。

self.labels = list(labels)self.shapes = np.array(shapes, dtype=np.float64)self.img_files = list(cache.keys()) # updateself.label_files = img2label_paths(cache.keys()) # updateif single_cls:for x in self.labels:x[:, 0] = 0n = len(shapes) # number of imagesbi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch indexnb = bi[-1] + 1 # number of batchesself.batch = bi # batch index of imageself.n = nself.indices = range(n)# Rectangular Training

Rectangular Training： 因yolov5在經過網絡層后，特征圖縮放為原來的1/32，所以輸入yolov5模型的分辨率必需能夠被32整除，例如當原始圖像是1280x720時，先等比縮放到640x360， 此時圖像的高為360，不是32的倍數，則將高填充到到384(比360大，且離360最近的能被32整除的數,計算方法為360/32向上取整，得到12，再用12*32即得到384)， 這樣高就是32地倍數，且填充的像素最少。此時高的兩邊需要分別填充(384-360)/2 = 12個像素。

if self.rect:# Sort by aspect ratios = self.shapes # whar = s[:, 1] / s[:, 0] # aspect ratioirect = ar.argsort()self.img_files = [self.img_files[i] for i in irect]self.label_files = [self.label_files[i] for i in irect]self.labels = [self.labels[i] for i in irect]self.shapes = s[irect] # whar = ar[irect]

以上操作是將數據按照圖像的寬高比排序。

# Set training image shapesshapes = [[1, 1]] * nbfor i in range(nb):ari = ar[bi == i]mini, maxi = ari.min(), ari.max()if maxi < 1:shapes[i] = [maxi, 1]elif mini > 1:shapes[i] = [1, 1 / mini]self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride

按照寬高比排序后，將寬高比接近的數據放到同一個batch。 因同一個batch里的數據，在輸入網絡時必需有相同的寬和高，這樣將寬高比接近的數據放到同一個batch，則batch內的數據，填充的無效像素就是最少的。
self.batch_shapes里面存放的就是每個batch最終輸入網絡的圖像的shape，這個shape是已經經過計算補充了無效像素的shape。

# Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)self.imgs, self.img_npy = [None] * n, [None] * nif cache_images:if cache_images == 'disk':self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + '_npy')self.img_npy = [self.im_cache_dir / Path(f).with_suffix('.npy').name for f in self.img_files]self.im_cache_dir.mkdir(parents=True, exist_ok=True)gb = 0 # Gigabytes of cached imagesself.img_hw0, self.img_hw = [None] * n, [None] * nresults = ThreadPool(NUM_THREADS).imap(lambda x: load_image(*x), zip(repeat(self), range(n)))pbar = tqdm(enumerate(results), total=n)for i, x in pbar:if cache_images == 'disk':if not self.img_npy[i].exists():np.save(self.img_npy[i].as_posix(), x[0])gb += self.img_npy[i].stat().st_sizeelse:self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)gb += self.imgs[i].nbytespbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})'pbar.close()

將數據集中所有圖像讀取進來，并等比縮放，保存到self.imgs里，當設置了本地保存時，保存為本地.npy文件

__getitem__函數解讀

想要使用pdb斷點到__getitem__函數，必需將dataloader的num_workers設置為0，如啟動訓練時將–workers參數設置為0即可。

該函數在train.py的292行 for i, (imgs, targets, paths, _) in pbar 調用，pbar = enumerate(train_loader)

def __getitem__(self, index):

輸入的index即是for i, (imgs, targets, paths, _) in pbar的值， 范圍是0 ~ batch num-1

index = self.indices[index] # linear, shuffled, or image_weightshyp = self.hypmosaic = self.mosaic and random.random() < hyp['mosaic']if mosaic:# Load mosaicimg, labels = load_mosaic(self, index)

訓練時采用mosaic數據增強方式，load_mosaic將隨機選取4張圖片組合成一張圖片，輸出的img size為self.img_size*self.img_size，如 640*640，參見load_mosaic函數解讀

shapes = None# MixUp augmentationif random.random() < hyp['mixup']:img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1)))

重新選取4張圖像mosaic增強，將mosaic增強后的圖像與之前mosaic增強后的數據進行mixup數據增強。

else:# Load imageimg, (h0, w0), (h, w) = load_image(self, index)

在推理時，單張圖片推理，yolov5使用自適應圖片縮放的方式，減少填充的黑邊以減少計算量。
# Letterbox shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape 在__init__函數講解中有介紹過self.batch_shapes里面存放的是每個batch最終輸入網絡的圖像的shape，這個shape是已經經過計算補充了無效像素的shape。

img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)

letterbox具體實現自適應縮放過程，scaleup控制是否向上縮放(即放大圖片)參見letterbox函數解讀。

shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescalinglabels = self.labels[index].copy()if labels.size: # normalized xywh to pixel xyxy formatlabels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])

img縮放與補黑邊后，label的框需要適配，xywhn2xyxy是將標注的label中歸一化的xywh中心點+寬高 -> xyxy左上角+右下角坐標，再加上補邊的偏移。

if self.augment:img, labels = random_perspective(img, labels,degrees=hyp['degrees'],translate=hyp['translate'],scale=hyp['scale'],shear=hyp['shear'],perspective=hyp['perspective'])

隨機裁剪縮放等數據增強。

nl = len(labels) # number of labelsif nl:labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)

xyxy左上角+右下角坐標 -> 歸一化的xywh中心點+寬高， clip規范xyxy坐標在圖片寬高內。

if self.augment:# Albumentationsimg, labels = self.albumentations(img, labels)

對圖像進行模糊等數據增強操作

nl = len(labels) # update after albumentations# HSV color-spaceaugment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])

augment_hsv函數將圖像轉到HSV空間進行數據增強，再轉回BGR

# Flip up-downif random.random() < hyp['flipud']:img = np.flipud(img)if nl:labels[:, 2] = 1 - labels[:, 2]# Flip left-rightif random.random() < hyp['fliplr']:img = np.fliplr(img)if nl:labels[:, 1] = 1 - labels[:, 1]# Cutouts# labels = cutout(img, labels, p=0.5)labels_out = torch.zeros((nl, 6))if nl:labels_out[:, 1:] = torch.from_numpy(labels)# Convertimg = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGBimg = np.ascontiguousarray(img)return torch.from_numpy(img), labels_out, self.img_files[index], shapes

load_mosaic函數解讀

mosaic數據增強是隨機選取4張圖片，然后拼合成一張圖片，整體步驟如下（以網絡的輸入img_size=640為例）：
1、 生成一張1280*1280的背景圖片，1280*1280的圖片是可以容納4張640*640的圖片的。
2、在1280*1280的背景圖片的中心附近區域(圖片中心點的上下左右1/4區域，見下圖中的黃色虛線框)隨機選取一個點，這個點作為4張圖片的接合點,見下圖中的紅色點。
3、隨機選取4張圖片，以上一步隨機選取的點為接合點，將4張圖片排列在背景圖片上作為前景。

4、裁剪掉超出背景圖片邊界的前景圖，得到1280*1280的圖像(代碼實現中是在每一張圖貼到背景圖的時候就去裁剪)。

5、將1280*1280的圖像縮放，隨機裁剪等方法得到640*640的圖像。

# loads images in a 4-mosaiclabels4, segments4 = [], [] s = self.img_size yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border]

圖片中心點的上下左右1/4區域隨機選取一個點，這個點作為4張圖片的接合點。

# mosaic center x, y indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices

隨機選取3張圖片，加上當前圖片，就是4張圖片

for i, index in enumerate(indices):# Load imageimg, _, (h, w) = load_image(self, index)

4張圖片分別貼到背景圖的top left，top right，bottom left，bottom right

# place img in img4if i == 0: # top leftimg4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)

第一張圖的時候創建背景圖

# base image with 4 tilesx1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)

x1a, y1a, x2a, y2a表示要貼在背景圖的位置

x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)

x1b, y1b, x2b, y2b表示裁剪后的前景圖坐標

elif i == 1: # top rightx1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), ycx1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), helif i == 2: # bottom leftx1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)elif i == 3: # bottom rightx1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]

將裁剪后的前景圖貼在背景圖上

padw = x1a - x1bpadh = y1a - y1b# Labelslabels, segments = self.labels[index].copy(), self.segments[index].copy()if labels.size:labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format

歸一化的xywh中心點+寬高 -> xyxy左上角+右下角坐標

segments = [xyn2xy(x, w, h, padw, padh) for x in segments]labels4.append(labels)segments4.extend(segments)# Concat/clip labels labels4 = np.concatenate(labels4, 0) for x in (labels4[:, 1:], *segments4):np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective() # img4, labels4 = replicate(img4, labels4) # replicate

以上都是在根據mosaic數據增強處理label，讓label能對應上拼合后的4合1圖像。

# Augment img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])

copy paste數據增強，在分割中使用，檢測中未使用，copy paste數據增強方式如下圖：

img4, labels4 = random_perspective(img4, labels4, segments4,degrees=self.hyp['degrees'],translate=self.hyp['translate'],scale=self.hyp['scale'],shear=self.hyp['shear'],perspective=self.hyp['perspective'],border=self.mosaic_border) # border to remove

縮放，隨機裁剪等，最終得到self.img_size*self.img_size大小的圖像

return img4, labels4

letterbox函數解讀（utils/augmentations.py）

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):

letterbox函數將輸入圖像im等比縮放到new_shape大小， 不夠的地方補充黑邊(或灰邊)

# Resize and pad image while meeting stride-multiple constraints shape = im.shape[:2] # current shape [height, width] if isinstance(new_shape, int):new_shape = (new_shape, new_shape)# Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better val mAP)r = min(r, 1.0)# Compute padding ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangledw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding

auto表示是否自動補齊到32的整數倍

elif scaleFill: # stretchdw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios

scaleFill表示不采用自適應縮放，直接resize到目標shape，無需補邊

dw /= 2 # divide padding into 2 sides dh /= 2if shape[::-1] != new_unpad: # resize

當原始圖像與目標shape不一致的時候，縮放

im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)

.

top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1))

dw與dh在上面除以2之前可能是個奇數，那么補邊的時候，兩側的邊補充的像素數應該是相差1的，上面兩步是在處理此種情況。

im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border

兩側補邊操作

return im, ratio, (dw, dh)

二、模型構建代碼解讀

這里只介紹檢測頭的構建，模型主干的構建沒有什么大的難點。

Detect類代碼解讀(models/yolo.py)

Detect類負責yolov5的3個檢測頭層的構建， 對應模型配置文件models/yolov5s.yaml中的最后一層，實現上只有一個卷積層，卷積核為1x1， 輸入是P3， P4， P5層。 Detect層位于下圖中的紅框處

__init__函數解讀

def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layersuper().__init__()self.nc = nc # number of classes

對于coco數據集來說， nc = 80

self.no = nc + 5 # number of outputs per anchor

需要預測的box的維度， xywh+正樣本置信度+80個類別每個類別的概率。

self.nl = len(anchors) # number of detection layersself.na = len(anchors[0]) // 2 # number of anchorsself.grid = [torch.zeros(1)] * self.nl # init grida = torch.tensor(anchors).float().view(self.nl, -1, 2)self.register_buffer('anchors', a) # shape(nl,na,2)self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2)self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv

ch = 3，表示輸入detect層的channel數， detect層的實現就是這個卷積核為1x1， 輸出channel為self.no*self.na = 85*3 = 255的卷積層。

self.inplace = inplace # use in-place ops (e.g. slice assignment)

forward函數解讀

def forward(self, x):

x是一個list，是detect層的輸入，list的長度為3，shape分別是(n, 128, 80, 80)， (n, 256, 40, 40), (n, 512, 20, 20)

# x = x.copy() # for profilingz = [] # inference output

以下代碼解讀均以i = 0來說明:

for i in range(self.nl):x[i] = self.m[i](x[i]) # conv

將執行卷積操作，執行完后，x[0]的shape為(n, 256, 80, 80)

bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

將x[0]的shape從(n, 256, 80, 80)轉換為(n, 3, 80, 80, 85)， 訓練時，直接返回這個x

if not self.training: # inferenceif self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

預測的box的中心點x和y是相對位置， 在推理的時候，需要映射到原圖上，因此需要先加上grid的坐標，再乘以stride映射回原圖。這里就是先把grid坐標計算出來。grid其實就是特征圖的每一個點的坐標。

y = x[i].sigmoid()if self.inplace:y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy

這里將預測框的中心點坐標轉化到-0.5~1.5范圍，然后將中心點映射回原圖。至于為什么要將預測框的中心點坐標轉化到-0.5~1.5范圍，可以參考后面ComputeLoss代碼解讀的 __call__函數解讀部分。

y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh

這里將預測框的寬和高轉化到0~4范圍，然后乘以預設anchor的寬和高(預設的anchor是基于原圖的，乘以anchor后可以將寬和高映射到原圖)，這里和計算loss時也是一致的。

else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xywh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # why = torch.cat((xy, wh, y[..., 4:]), -1)z.append(y.view(bs, -1, self.no))

推理時，將每一個檢測頭的預測結果存放在z中，z[0]的shape為(n, 19200, 85)， 其中19200 = 3 * 80 * 80, 表示第一層detect檢測頭所有grid預測的box信息。

return x if self.training else (torch.cat(z, 1), x)

三、loss計算代碼解讀

ComputeLoss代碼解讀(utils/loss.py)

__init__函數解讀

該函數在train.py的258行compute_loss = ComputeLoss(model)調用

def __init__(self, model, autobalance=False): super(ComputeLoss, self).__init__() self.sort_obj_iou = False device = next(model.parameters()).device # get model device h = model.hyp # hyperparameters# Define criteria BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets# Focal loss g = h['fl_gamma'] # focal loss gamma if g > 0:BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module

這里det返回的是檢測層，對應的是models/yolo.py->class Detect->__init__->self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)構建的檢測頭。詳見class Detect章節

返回的det層打印如下：

(Pdb) p det Detect((m): ModuleList((0): Conv2d(128, 255, kernel_size=(1, 1), stride=(1, 1))(1): Conv2d(256, 255, kernel_size=(1, 1), stride=(1, 1))(2): Conv2d(512, 255, kernel_size=(1, 1), stride=(1, 1))) )

.

self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7

det.nl是3，表示檢測頭的數目，self.balance的結果還是[4.0, 1.0, 0.4]，標識三個檢測頭對應輸出的損失系數。

self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index

det.stride是[ 8., 16., 32.]， self.ssi表示stride為16的索引，當autobalance為true時，self.ssi為1.

self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance for k in 'na', 'nc', 'nl', 'anchors':setattr(self, k, getattr(det, k))

__call__函數解讀

該函數在train.py的318行loss, loss_items = compute_loss(pred, targets.to(device))調用。
函數定義

def __call__(self, p, targets):

輸入1：
p， 是每個預測頭輸出的結果，

[p[0].shape： torch.Size([16, 3, 80, 80, 85])p[1].shape： torch.Size([16, 3, 40, 40, 85])p[2].shape： torch.Size([16, 3, 20, 20, 85]) ]

p[0]的每個維度解釋：
16： batch size
3： anchor box數量
80/40/20： 3個檢測頭特征圖大小
85： coco數據集80個類別+4(x,y,w,h)+1(是否為前景)

輸入2:
targets: gt box信息，維度是(n, 6)，其中n是整個batch的圖片里gt box的數量，以下都以gt box數量為190來舉例。 6的每一個維度為(圖片在batch中的索引， 目標類別， x, y, w, h)

device = targets.device lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets

tcls的shape為(3, 808), 表示3個檢測頭對應的gt box的類別。

tbox的shape為(3, ([808, 4]))， 表示3個檢測頭對應的gt box的xywh， 其中x和y已經減去了預測方格的整數坐標，
比如原始的gt box的中心坐標是(51.7, 44.8)，則該gt box由方格(51, 44)，以及離中心點最近的兩個方格(51, 45)和(52, 44)來預測(見build_targets函數里的解析),
換句話說這三個方格預測的gt box是同一個，其中心點是(51.7, 44.8)，但tbox保存這三個方格預測的gt box的xy時，保存的是針對這三個方格的偏移量,
分別是：
(51.7 - 51 = 0.7, 44.8 - 44 = 0.8)
(51.7 - 51 = 0.7, 44.8 - 45 = -0.2)
(51.7 - 52 = -0.3, 44.8 - 44 = 0.8)

indices的shape為(3, ([808], [808], [808], [808])), 4個808分別表示每個gt box(包括偏移后的gt box)在batch中的image index， anchor index， 預測該gt box的網格y坐標， 預測該gt box的網格x坐標。

anchors的shape為(3, ([808, 2]))， 表示每個檢測頭對應的808個gt box所對應的anchor。

以下都以第一個檢測頭對應的808個gt box講解，其他檢測頭的gt box個數不同，但邏輯相同

for i, pi in enumerate(p): # layer index, layer predictionsb, a, gj, gi = indices[i] # image, anchor, gridy, gridxtobj = torch.zeros_like(pi[..., 0], device=device) # target objn = b.shape[0] # number of targetsif n:ps = pi[b, a, gj, gi] # prediction subset corresponding to targets

將808個gt box對應的預測框選取出來。 ps的shape為(808, 85)。

# Regressionpxy = ps[:, :2].sigmoid() * 2. - 0.5

參考

https://github.com/ultralytics/yolov5/issues/1585#issuecomment-739060912?github.com/ultralytics/yolov5/issues/1585#issuecomment-739060912

將預測的中心點坐標變換到-0.5到1.5之間，如下圖所示：

因為像素是沒有小數的，但預測出來的坐標都是帶小數位的，所以這里把一個像素當成一個方格來對待，如上圖所示，將中心坐標變換到-0.5到1.5之間就相當于預測的中心點范圍限制在圖中的綠色框中，也就是說(51, 44)這個方格可預測的目標中心點范圍是(50.5, 43.5)到(52.5, 45.5)之間。

pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]

參考

https://github.com/ultralytics/yolov5/issues/471#issuecomment-662009779?github.com/ultralytics/yolov5/issues/471#issuecomment-662009779

大概意思解釋下：
ps[:, 2:4].sigmoid() * 2) ** 2的范圍是0~4， 再乘以anchors[i]， 表示把預測框的寬和高限制在4倍的anchors內，這是為了解決yolov3和yolov4對預測框寬高無任何約束的問題，這個4和默認的超參數anchor_t是相等的，也是有關聯的，調整時建議一起調整。

pbox = torch.cat((pxy, pwh), 1) # predicted boxiou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)

此處就是計算808個預測框與gt box的giou了， 得到的iou的shape是(808)。

lbox += (1.0 - iou).mean() # iou loss# Objectnessscore_iou = iou.detach().clamp(0).type(tobj.dtype)

detach函數使得iou不可反向傳播， clamp將小于0的iou裁剪為0

if self.sort_obj_iou:sort_id = torch.argsort(score_iou)

torch.argsort返回的是排序后的score_iou中的元素在原始score_iou中的位置。

b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id]

得到根據iou從小到大排序的image index, anchor index, gridy, gridx, iou。

tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # giou ratio

tobj[b, a, gj, gi]的shape為808，表示選擇有對應gt box的808個預測框的位置賦值，其他位置的值為初始值0。 tobj的shape為(16, 3, 80, 80), 表示每個grid的預測框與gt box的iou。

# Classificationif self.nc > 1: # cls loss (only if multiple classes)t = torch.full_like(ps[:, 5:], self.cn, device=device) # targetst[range(n), tcls[i]] = self.cplcls += self.BCEcls(ps[:, 5:], t) # BCE

ps[:, 5:]取808個預測框信息的第6-85個數據，即目標是每個類別的概率。
http://self.cn和self.cp分別是標簽平滑的負樣本平滑標簽和正樣本平滑標簽， 參考https://blog.csdn.net/qq_38253797/article/details/116228065
tcls的shape為(3, 808), 表示3個檢測頭對應的gt box的類別，tcls[i]的shape為(808)， 表示808個gt box的類別， 取值為0~79。
t[range(n), tcls[i]]的shape是(808, 80)， t的每一行里面有一個值是self.cp(即正樣本平滑標簽)，其他值是http://self.cn(即負樣本平滑標簽)

# Append targets to text file# with open('targets.txt', 'a') as file:# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]obji = self.BCEobj(pi[..., 4], tobj)

tobj的shape為(16, 3, 80, 80), 表示每個grid的預測框與gt box的iou。
pi[…, 4]取每個grid預測的object為正樣本的概率。
將每個grid預測框和gt box的iou與每個grid預測的object為正樣本的概率計算BCE loss，做為object loss。

lobj += obji * self.balance[i] # obj lossif self.autobalance:self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item() if self.autobalance:self.balance = [x / self.balance[self.ssi] for x in self.balance] lbox *= self.hyp['box'] lobj *= self.hyp['obj'] lcls *= self.hyp['cls'] bs = tobj.shape[0] # batch size return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()

build_targets函數解讀

函數定義

def build_targets(self, p, targets):

輸入1：p， 是每個預測頭輸出的結果，

[p[0].shape： torch.Size([16, 3, 80, 80, 85])p[1].shape： torch.Size([16, 3, 40, 40, 85])p[2].shape： torch.Size([16, 3, 20, 20, 85]) ]

p[0]的每個維度解釋：
16： batch size
3： anchor box數量
80/40/20： 3個檢測頭特征圖大小
85： coco數據集80個類別+4(x,y,w,h)+1(是否為前景)

輸入2: targets: gt box信息，維度是(n, 6)，其中n是整個batch的圖片里gt box的數量，以下都以gt box數量為190來舉例。 6的每一個維度為(圖片在batch中的索引， 目標類別， x, y, w, h)

該函數主要是處理gt box，先介紹一下gt box的整體處理策略：
1、將gt box復制3份，原因是有三種長寬的anchor， 每種anchor都有gt box與其對應，也就是在篩選之前，一個gt box有三種anchor與其對應。
2、過濾掉gt box的w和h與anchor的w和h的比值大于設置的超參數anchor_t的gt box。
3、剩余的gt box，每個gt box使用至少三個方格來預測，一個是gt box中心點所在方格，另兩個是中心點離的最近的兩個方格，如下圖：

如果gt box的中心坐標是(51.7, 44.8)，則由gt box中心點所在方格(51, 44)，以及離中心點最近的兩個方格(51, 45)和(52, 44)來預測此gt box。

# Build targets for compute_loss(), input targets(image,class,x,y,w,h) na, nt = self.na, targets.shape[0] # number of anchors： 3, gt box的數量：190 tcls, tbox, indices, anch = [], [], [], [] gain = torch.ones(7, device=targets.device) # normalized to gridspace gain

gain = [1, 1, 1, 1, 1, 1, 1], 是7個數，前6個數對應targets的第二維度6，先不用關注用途，后面會講到。

ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)

anchor的索引，shape為(3, gt box的數量)， 3行里，第一行全是0， 第2行全是1， 第三行全是2，表示每個gt box都對應到3個anchor上。

targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices

下圖顯示了處理過后的targets內容，其shape為(3, 190, 7), 紅色框中是原始的targets， 藍色框是給每個gt box加上索引，表示對應著哪種anchor。 然后將gt box重復3遍，對應著三種anchor.

g = 0.5 # bias off = torch.tensor([[0, 0],[1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm], device=targets.device).float() * g # offsets

以下都以第一個檢測頭為例講解，其他檢測頭的邏輯相同

for i in range(self.nl): // 針對每一個檢測頭anchors = self.anchors[i]

i=0時， anchors值為：

(Pdb) p anchors tensor([[1.25000, 1.62500],[2.00000, 3.75000],[4.12500, 2.87500]], device='cuda:0')

.

gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain

gain的值是[ 1., 1., 80., 80., 80., 80., 1.]， 80的位置對應著gt box信息的xywh。

# Match targets to anchors t = targets * gain

targets里的xywh是歸一化到0 ~ 1之間的， 乘以gain之后，將targets的xywh映射到檢測頭的特征圖大小上。

if nt:# Matchesr = t[:, :, 4:6] / anchors[:, None] # wh ratio

這里的r的shape為[3, 190, 2]， 2分別表示gt box的w和h與anchor的w和h的比值。

j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare

當gt box的w和h與anchor的w和h的比值比設置的超參數anchor_t大時，則此gt box去除，這一步得到的j的shape為(3, 190)， 里面的值均為true或false， 表示每一個gt box是否將要過濾掉。

# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))t = t[j] # filter

此步就是去除gt box的w和h與anchor的w和h的比值大于anchor_t的gt box，得到的t的shape為(271, 7)， 這個271是本人試驗的結果，不同的數據集會不一樣，可以看到過濾前，3個anchor總的gt box是3 * 190 = 570個，過濾后還剩余271個。過濾的原因是和anchor的寬高差別較大的gt box，是非常難預測的，不適合用來訓練。

# Offsetsgxy = t[:, 2:4] # grid xy

取出過濾后的gt box的中心點浮點型的坐標。

gxi = gain[[2, 3]] - gxy # inverse

將以圖像左上角為原點的坐標變換為以圖像右下角為原點的坐標。

j, k = ((gxy % 1. < g) & (gxy > 1.)).T

以圖像左上角為原點的坐標，取中心點的小數部分，小數部分小于0.5的為ture，大于0.5的為false。 j和k的shape都是(271)，true的位置分別表示靠近方格左邊的gt box和靠近方格上方的gt box。

l, m = ((gxi % 1. < g) & (gxi > 1.)).T

以圖像右下角為原點的坐標，取中心點的小數部分，小數部分小于0.5的為ture，大于0.5的為false。 l和m的shape都是(271)，true的位置分別表示靠近方格右邊的gt box和靠近方格下方的gt box。
j和l的值是剛好相反的，k和m的值也是剛好相反的。

j = torch.stack((torch.ones_like(j), j, k, l, m))

將j, k, l, m組合成一個tensor，另外還增加了一個全為true的維度。組合之后，j的shape為(5, 271)

t = t.repeat((5, 1, 1))[j]

t之前的shape為(271, 7)， 這里將t復制5個，然后使用j來過濾，
第一個t是保留所有的gt box，因為上一步里面增加了一個全為true的維度，
第二個t保留了靠近方格左邊的gt box，
第三個t保留了靠近方格上方的gt box，
第四個t保留了靠近方格右邊的gt box，
第五個t保留了靠近方格下邊的gt box，
過濾后，t的shape為(808, 7), 表示保留下來的所有的gt box。

offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]

offsets的shape為(808, 2), 表示保留下來的808個gt box的x, y對應的偏移，
第一個t保留所有的gt box偏移量為[0, 0], 即不做偏移
第二個t保留的靠近方格左邊的gt box，偏移為[0.5, 0]，即向左偏移0.5(后面代碼是用gxy - offsets，所以正0.5表示向左偏移)，則偏移到左邊方格，表示用左邊的方格來預測
第三個t保留的靠近方格上方的gt box，偏移為[0, 0.5]，即向上偏移0.5，則偏移到上邊方格，表示用上邊的方格來預測
第四個t保留的靠近方格右邊的gt box，偏移為[-0.5, 0]，即向右偏移0.5，則偏移到右邊方格，表示用右邊的方格來預測
第五個t保留的靠近方格下邊的gt box，偏移為[0, 0.5]，即向下偏移0.5，則偏移到下邊方格，表示用下邊的方格來預測

一個gt box的中心點x坐標要么是靠近方格左邊，要么是靠近方格右邊，y坐標要么是靠近方格上邊，要么是靠近方格下邊，所以一個gt box在以上五個t里面，會有三個t是true。
也即一個gt box有三個方格來預測，一個是中心點所在方格，另兩個是離的最近的兩個方格。而yolov3只使用中心點所在的方格預測，這是與yolov3的區別。

else:t = targets[0]offsets = 0# Define b, c = t[:, :2].long().T # image, class gxy = t[:, 2:4] # grid xy gwh = t[:, 4:6] # grid wh gij = (gxy - offsets).long()

將中心點偏移到相鄰最近的方格里，然后向下取整， gij的shape為(808, 2)

gi, gj = gij.T # grid xy indices# Append a = t[:, 6].long() # anchor indices indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices tbox.append(torch.cat((gxy - gij, gwh), 1)) # box anch.append(anchors[a]) # anchors tcls.append(c) # classreturn tcls, tbox, indices, anch

總結

以上是生活随笔為你收集整理的YoloV5代码详细解读的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。