项目结构
这个是我做yolo项目的结构，总结一下，其他都中规中矩，比较麻烦的是数据的转换和loss的计算，不像我之前的项目，感觉自己一个人写不出来，主要感觉是太繁琐了，所以找了别人的github项目中的部分内容借鉴了一下
特别注意，你在使用这个项目之前要把数据集,class和anchor的txt准备好。
config
from model import YoloBody
from config import parser
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from torch.utils.data import DataLoader
from datalist import YoloDataset,yolo_dataset_collate
import numpy as np
from utils import YOLOLoss
the_loss=1000
class train(object):
def __init__(self):
self.args=parser.parse_args()
print(f"-----------{self.args.project_name}-------------")
use_cuda=self.args.use_cuda and torch.cuda.is_available()
if use_cuda:
torch.cuda.manual_seed(self.args.seed)
torch.cuda.manual_seed_all(self.args.seed)
else:
torch.manual_seed(self.args.seed)
self.device=torch.device('cuda' if use_cuda else 'cpu')
train_kwargs={'num_workers':0,'pin_memory':True} if use_cuda else {}
test_kwargs={'num_workers':0,'pin_memory':False} if use_cuda else {}
'''
构造DataLoader
'''
self.annotations=self.args.data_path
with open(self.annotations) as f:
self.lines=f.readlines()
np.random.seed(self.args.seed)
np.random.shuffle(self.lines)
np.random.seed(None)
self.num_val=int(len(self.lines)*self.args.val_num)
self.num_train=len(self.lines)-self.num_val
self.train_dataset=YoloDataset(self.lines[:self.num_train],self.args.image_size)
self.test_dataset=YoloDataset(self.lines[self.num_train:],self.args.image_size)
self.train_dataloader=DataLoader(self.train_dataset,batch_size=self.args.train_batchsize,shuffle=True,**train_kwargs,collate_fn=yolo_dataset_collate)
self.test_dataloader=DataLoader(self.test_dataset,batch_size=self.args.test_batchsize,shuffle=False,**test_kwargs,collate_fn=yolo_dataset_collate)
'''
构造model
'''
self.num_classes=self.get_classes(self.args.classes_path)
self.anchors=self.get_anchors(self.args.anchor_path)
self.model=YoloBody(len(self.anchors[0]),len(self.num_classes))
if use_cuda:
self.model=torch.nn.DataParallel(self.model,device_ids=range(torch.cuda.device_count()))
cudnn.benchmark=True
if self.args.resume:
try:
print("load the weight from pretrained-weight file")
model_dict = self.model.state_dict()
pretrained_dict = torch.load(self.args.pretrained_weight, map_location=self.device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
self.model.load_state_dict(model_dict)
print("Finished to load the weight")
except:
print("can not load weight \n train the model from stratch")
self.model.apply(self.weights_init)
'''
构造loss目标函数
选择优化器
学习率变化选择
'''
self.criterion=self.create_loss()
self.optimizer=optim.Adam(self.model.parameters(),self.args.lr)
self.scheduler=optim.lr_scheduler.CosineAnnealingLR(self.optimizer,T_max=5,eta_min=1e-5)
for epoch in range(1,self.args.epoches):
self.train(epoch)
if epoch %1==0:
self.test(epoch)
torch.cuda.empty_cache()
print("finish model training")
def train(self,epoch):
self.model.train()
losses = []
pbar=tqdm(self.train_dataloader,desc=f'Train Epoch{epoch}/{self.args.epoches}')
for data, target in pbar:
total_loss = 0
data,target=torch.from_numpy(data).type(torch.FloatTensor), [torch.from_numpy(ann).type(torch.FloatTensor) for ann in target]
data,target=data.to(self.device),target
self.optimizer.zero_grad()
outputs=self.model(data)
for i in range(3):
loss_item=self.criterion[i](outputs[i],target,self.device)[0]
total_loss+=loss_item
losses.append(loss_item.item())
loss=sum(losses)
total_loss.backward()
self.optimizer.step()
pbar.set_description(
f'Train Epoch:{epoch}/{self.args.epoches} train_loss:{round(np.mean(losses),4)}'
)
self.scheduler.step()
def test(self,epoch):
self.model.eval()
losses=[]
pbar=tqdm(self.test_dataloader,desc=f'Test Epoch{epoch}/{self.args.epoches}')
with torch.no_grad():
for data,target in pbar:
data, target = torch.from_numpy(data).type(torch.FloatTensor), [
torch.from_numpy(ann).type(torch.FloatTensor) for ann in target]
data,target=data.to(self.device),target
outputs=self.model(data)
for i in range(3):
loss_item=self.criterion[i](outputs[i],target)[0]
losses.append(loss_item.item())
pbar.set_description(
f'【Test】 Epoch:{epoch}/{self.args.epoches} test_loss:{round(np.mean(losses),4)}'
)
global the_loss
if np.mean(losses)>the_loss:
the_loss=np.mean(losses)
torch.save({
'epoch': epoch,
'model_state_dict': self.model.state_dict(),
"optimizer_state_dict": self.optimizer.state_dict(),
}
,'./weights/'+f'Epoch_'+str(epoch)+'_Loss_'+str(the_loss)+'.pth'
)
def get_classes(self,classes_path):
with open(classes_path) as f:
class_names=f.readlines()
class_names=[c.strip() for c in class_names]
return class_names
def get_anchors(self,anchors_path):
with open(anchors_path) as f:
anchors=f.readline()
anchors=[float(x) for x in anchors.split(',')]
return np.array(anchors).reshape(-1,3,2)[::-1,:,:]
#因为这里的loss比较复杂
def create_loss(self):
yolo_losses=[]
for i in range(3):
yolo_losses.append(YOLOLoss(np.reshape(self.anchors,[-1,2]),len(self.num_classes),
(416,416)))
return yolo_losses
train=train()
datalist
from random import shuffle
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
class YoloDataset(Dataset):
def __init__(self, train_lines, image_size):
super(YoloDataset, self).__init__()
self.train_lines = train_lines
self.image_size = image_size
def __len__(self):
return len(self.train_lines)
def __getitem__(self, index):
if index == 0:
shuffle(self.train_lines)
index = index % len(self.train_lines)
line = self.train_lines[index].split('.jpg')
line_path = line[0] + '.jpg'
line_label = line[-1].strip().split()
image = Image.open("E:/Datasets/mosaic_blood/" + line_path.split('/')[-1])
iw, ih = image.size
h, w = self.image_size
box = []
line_label = list(map(int, line_label))
for i in range(0, len(line_label), 5):
b = line_label[i:i + 5]
box.append(b)
box = np.array(box)
image = image.resize((w, h), Image.BICUBIC)
box_data = np.zeros((len(box), 5))
if len(box) > 0:
np.random.shuffle(box)
box[:,0:2][box[:, 0:2] < 0] = 0
box[:, 2][box[:, 2] > w] = w
box[:, 3][box[:, 3] > h] = h
box_w = box[:, 2] - box[:, 0]
box_h = box[:, 3] - box[:, 1]
box = box[np.logical_and(box_w > 1, box_h > 1)] # 保留有效框
box_data = np.zeros((len(box), 5))
box_data[:len(box)] = box
if len(box) == 0:
return image, []
if (box_data[:, :4] > 0).any():
# 从坐标转换成0~1的百分比
boxes = np.array(box_data[:, :4], dtype=np.float32)
boxes[:, 0] = boxes[:, 0] / self.image_size[1]
boxes[:, 1] = boxes[:, 1] / self.image_size[0]
boxes[:, 2] = boxes[:, 2] / self.image_size[1]
boxes[:, 3] = boxes[:, 3] / self.image_size[0]
boxes = np.maximum(np.minimum(boxes, 1), 0)
boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
boxes[:, 0] = boxes[:, 0] + boxes[:, 2] / 2
boxes[:, 1] = boxes[:, 1] + boxes[:, 3] / 2
y = np.concatenate([boxes, box_data[:, -1:]], axis=-1)
image=np.array(image,dtype=np.float32)
tmp_img=np.transpose(image/255.0,(2,0,1))
tmp_targets=np.array(y,dtype=np.float32)
return tmp_img,tmp_targets
else:
return image, []
def yolo_dataset_collate(batch):
images=[]
bboxes=[]
for img,box in batch:
images.append(img)
bboxes.append(box)
images=np.array(images)
bboxes=np.array(bboxes)
return images,bboxes
model
import math
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
class Mish(nn.Module):
def __init__(self):
super(Mish, self).__init__()
def forward(self,x):
return x*torch.tanh(F.softplus(x))
class BasicConv(nn.Module):
def __init__(self,in_channels,out_channels,kernel_size,stride=1):
super(BasicConv, self).__init__()
self.conv=nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding=kernel_size//2,bias=False)
self.bn=nn.BatchNorm2d(out_channels)
self.activation=Mish()
def forward(self,x):
x=self.conv(x)
x=self.bn(x)
x=self.activation(x)
return x
class Resblock(nn.Module):
def __init__(self,channels,hidden_channels=None,residual_activation=nn.Identity()):
super(Resblock, self).__init__()
if hidden_channels is None:
hidden_channels=channels
self.block=nn.Sequential(
BasicConv(channels,hidden_channels,1),
BasicConv(hidden_channels,channels,3)
)
def forward(self,x):
return x+self.block(x)
class Resblock_body(nn.Module):
def __init__(self, in_channels, out_channels, num_blocks, first):
super(Resblock_body, self).__init__()
self.downsample_conv = BasicConv(in_channels, out_channels, 3, stride=2)
if first:
self.split_conv0 = BasicConv(out_channels, out_channels, 1)
self.split_conv1 = BasicConv(out_channels, out_channels, 1)
self.block_conv = nn.Sequential(
Resblock(channels=out_channels, hidden_channels=out_channels // 2),
BasicConv(out_channels, out_channels, 1)
)
self.concat_conv = BasicConv(out_channels * 2, out_channels, 1)
else:
self.split_conv0 = BasicConv(out_channels, out_channels // 2, 1)
self.split_conv1 = BasicConv(out_channels, out_channels // 2, 1)
self.block_conv = nn.Sequential(
*[Resblock(out_channels // 2) for _ in range(num_blocks)],
BasicConv(out_channels // 2, out_channels // 2, 1)
)
self.concat_conv = BasicConv(out_channels, out_channels, 1)
def forward(self, x):
x = self.downsample_conv(x)
x0 = self.split_conv0(x)
x1 = self.split_conv1(x)
x1 = self.block_conv(x1)
x = torch.cat([x1, x0], dim=1)
x = self.concat_conv(x)
return x
class CSPDarkNet(nn.Module):
def __init__(self, layers):
super(CSPDarkNet, self).__init__()
self.inplanes = 32
self.conv1 = BasicConv(3, self.inplanes, kernel_size=3, stride=1)
self.feature_channels = [64, 128, 256, 512, 1024]
self.stages = nn.ModuleList([
Resblock_body(self.inplanes, self.feature_channels[0], layers[0], first=True),
Resblock_body(self.feature_channels[0], self.feature_channels[1], layers[1], first=False),
Resblock_body(self.feature_channels[1], self.feature_channels[2], layers[2], first=False),
Resblock_body(self.feature_channels[2], self.feature_channels[3], layers[3], first=False),
Resblock_body(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)
])
self.num_features = 1
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def forward(self, x):
x = self.conv1(x)
x = self.stages[0](x)
x = self.stages[1](x)
out3 = self.stages[2](x)
out4 = self.stages[3](out3)
out5 = self.stages[4](out4)
return out3, out4, out5
def darknet53(pretrained, **kwargs):
model = CSPDarkNet([1, 2, 8, 8, 4])
if pretrained:
if isinstance(pretrained, str):
model.load_state_dict(torch.load(pretrained))
else:
raise Exception("darknet request a pretrained path. got [{}]".format(pretrained))
return model
def conv2d(filter_in, filter_out, kernel_size, stride=1):
pad = (kernel_size - 1) // 2 if kernel_size else 0
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)),
])
)
class SpatialPytamidPooling(nn.Module):
def __init__(self, pool_sizes=[5, 9, 13]):
super(SpatialPytamidPooling, self).__init__()
self.maxpools = nn.ModuleList([
nn.MaxPool2d(kernel_size=pool_size, stride=1, padding=pool_size // 2) for pool_size in pool_sizes
])
def forward(self, x):
features = [maxpool(x) for maxpool in self.maxpools[::-1]]
features = torch.cat(features + [x], dim=1)
return features
class Upsample(nn.Module):
def __init__(self, in_channels, out_channels):
super(Upsample, self).__init__()
self.upsample = nn.Sequential(
conv2d(in_channels, out_channels, 1),
nn.Upsample(scale_factor=2, mode='nearest')
)
def forward(self, x):
x = self.upsample(x)
return x
def make_three_conv(filter_list, in_filters):
m = nn.Sequential(
conv2d(in_filters, filter_list[0], 1),
conv2d(filter_list[0], filter_list[1], 3),
conv2d(filter_list[1], filter_list[0], 1)
)
return m
def make_five_conv(filter_list, in_filters):
m = nn.Sequential(
conv2d(in_filters, filter_list[0], 1),
conv2d(filter_list[0], filter_list[1], 3),
conv2d(filter_list[1], filter_list[0], 1),
conv2d(filter_list[0], filter_list[1], 3),
conv2d(filter_list[1], filter_list[0], 1)
)
return m
def yolo_head(filtes_list, in_filters):
m = nn.Sequential(
conv2d(in_filters, filtes_list[0], 3),
nn.Conv2d(filtes_list[0], filtes_list[1], 1)
)
return m
class YoloBody(nn.Module):
def __init__(self, num_anchors, num_classes):
super(YoloBody, self).__init__()
self.backbone = darknet53(None)
self.conv1 = make_three_conv([512, 1024], 1024)
self.SPP = SpatialPytamidPooling()
self.conv2 = make_three_conv([512, 1024], 2048)
self.upsample1 = Upsample(512, 256)
self.conv_for_P4 = conv2d(512, 256, 1)
self.make_five_conv1 = make_five_conv([256, 512], 512)
self.upsample2 = Upsample(256, 128)
self.conv_for_P3 = conv2d(256, 128, 1)
self.make_five_conv2 = make_five_conv([128, 256], 256)
# 4+1+num_classes
final_out_filter2 = num_anchors * (5 + num_classes)
self.yolo_head3 = yolo_head([256, final_out_filter2], 128)
self.down_sample1 = conv2d(128, 256, 3, stride=2)
self.make_five_conv3 = make_five_conv([256, 512], 512)
# 3*(5+num_classes)=3*(5+20)=3*(4+1+20)=75
final_out_filter1 = num_anchors * (5 + num_classes)
self.yolo_head2 = yolo_head([512, final_out_filter1], 256)
self.down_sample2 = conv2d(256, 512, 3, stride=2)
self.make_five_conv4 = make_five_conv([512, 1024], 1024)
# 3*(5+num_classes)=3*(5+20)=3*(4+1+20)=75
final_out_filter0 = num_anchors * (5 + num_classes)
self.yolo_head1 = yolo_head([1024, final_out_filter0], 512)
def forward(self, x):
"""
:param x: shape is (N, C_in, H_in, W_in)
:return: out0 shape is (N, num_anchors*(5+num_classes), H_in/8, W_in/8)
out1 shape is (N, num_anchors*(5+num_classes), H_in/16, W_in/16)
out2 shape is (N, num_anchors*(5+num_classes), H_in/32, W_in/32)
"""
# backbone
x2, x1, x0 = self.backbone(x)
P5 = self.conv1(x0)
P5 = self.SPP(P5)
P5 = self.conv2(P5)
P5_upsample = self.upsample1(P5)
P4 = self.conv_for_P4(x1)
P4 = torch.cat([P4, P5_upsample], axis=1)
P4 = self.make_five_conv1(P4)
P4_upsample = self.upsample2(P4)
P3 = self.conv_for_P3(x2)
P3 = torch.cat([P3, P4_upsample], axis=1)
P3 = self.make_five_conv2(P3)
P3_downsample = self.down_sample1(P3)
P4 = torch.cat([P3_downsample, P4], axis=1)
P4 = self.make_five_conv3(P4)
P4_downsample = self.down_sample2(P4)
P5 = torch.cat([P4_downsample, P5], axis=1)
P5 = self.make_five_conv4(P5)
out2 = self.yolo_head3(P3)
out1 = self.yolo_head2(P4)
out0 = self.yolo_head1(P5)
return out0, out1, out2
utils
import torch.nn as nn
import torch
import math
import numpy as np
def jaccard(_box_a, _box_b):
"""
:param _box_a: gt_box shape is (N, 4) x, y, w, h
:param _box_b: shape is (M, 4)
:return: inter / union shape is (N, M)
"""
# b1_x1, b1_x2, b1_y1, b1_y2 shape is N
b1_x1, b1_x2 = _box_a[:, 0] - _box_a[:, 2] / 2, _box_a[:, 0] + _box_a[:, 2] / 2 # b1_x1左上x坐标,b1_x2右下x坐标
b1_y1, b1_y2 = _box_a[:, 1] - _box_a[:, 3] / 2, _box_a[:, 1] + _box_a[:, 3] / 2 # b1_y1左上y坐标,b1_y2右下y坐标
# b2_x1, b2_x2, b2_y1, b2_y2 shape is M
b2_x1, b2_x2 = _box_b[:, 0] - _box_b[:, 2] / 2, _box_b[:, 0] + _box_b[:, 2] / 2 # b2_x1左上x坐标,b2_x2右下x坐标
b2_y1, b2_y2 = _box_b[:, 1] - _box_b[:, 3] / 2, _box_b[:, 1] + _box_b[:, 3] / 2 # b2_y1左上y坐标,b2_y2右下y坐标
box_a = torch.zeros_like(_box_a) # shape is (N, 4)
box_b = torch.zeros_like(_box_b) # shape is (M, 4)
box_a[:, 0], box_a[:, 1], box_a[:, 2], box_a[:, 3] = b1_x1, b1_y1, b1_x2, b1_y2 # box_a shape is (N, 4)
box_b[:, 0], box_b[:, 1], box_b[:, 2], box_b[:, 3] = b2_x1, b2_y1, b2_x2, b2_y2 # box_b shape is (M, 4)
# box_a存储左上右下的坐标 box_b存储左上右下的坐标
A = box_a.size(0) # A=N
B = box_b.size(0) # B=M
max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), # shape is (N, M, 2)
box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) # shape is (N, M, 2)
# max_xy shape is (N, M, 2)，min_xy shape is (N, M, 2)
min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
box_b[:, :2].unsqueeze(0).expand(A, B, 2))
inter = torch.clamp((max_xy - min_xy), min=0) # shape is (N, M, 2)
inter = inter[:, :, 0] * inter[:, :, 1] # shape is (N, M)
# 计算先验框和真实框各自的面积
area_a = ((box_a[:, 2] - box_a[:, 0]) *
(box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter) # shape is (N, M)
area_b = ((box_b[:, 2] - box_b[:, 0]) *
(box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter) # shape is (N, M)
# 求IOU
union = area_a + area_b - inter # shape is (N, M)
return inter / union # shape is (N, M)
def clip_by_tensor(t, t_min, t_max):
t = t.float()
result = (t >= t_min).float() * t + (t < t_min).float() * t_min
result = (result <= t_max).float() * result + (result > t_max).float() * t_max
return result
def bbox_iou(box1, box2, x1y1x2y2=True):
"""
计算IOU
:param box1: shape is (1, 4)
:param box2: shape is (M, 4)
:param x1y1x2y2:
:return: iou shape is M
"""
if not x1y1x2y2:
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
else:
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * \
torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)
iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)
return iou
def box_ciou(b1, b2):
"""
:param b1: shape is (N, 4)，调整以后的预测框中心与宽高
:param b2: shape is (N, 4), tx,ty,tw,th，标签框的值，这些值没做归一化
:return: ciou shape is N
"""
"""
:param b1: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh, 预测
:param b2: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh, 标签
:return: ciou: tensor, shape=(batch, feat_w, feat_h, anchor_num, 1)
"""
# 求出预测框左上角右下角
b1_xy = b1[..., :2]
b1_wh = b1[..., 2:4]
b1_wh_half = b1_wh / 2.
b1_mins = b1_xy - b1_wh_half
b1_maxes = b1_xy + b1_wh_half
# 求出真实框左上角右下角
b2_xy = b2[..., :2]
b2_wh = b2[..., 2:4]
b2_wh_half = b2_wh / 2.
b2_mins = b2_xy - b2_wh_half
b2_maxes = b2_xy + b2_wh_half
# 求真实框和预测框所有的iou
intersect_mins = torch.max(b1_mins, b2_mins)
intersect_maxes = torch.min(b1_maxes, b2_maxes)
intersect_wh = torch.max(intersect_maxes - intersect_mins, torch.zeros_like(intersect_maxes))
intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
b1_area = b1_wh[..., 0] * b1_wh[..., 1]
b2_area = b2_wh[..., 0] * b2_wh[..., 1]
union_area = b1_area + b2_area - intersect_area
iou = intersect_area / torch.clamp(union_area, min=1e-6)
# 计算中心的差距
center_distance = torch.sum(torch.pow((b1_xy - b2_xy), 2), axis=-1) # p(Actr, Bctr)^2
# 找到包裹两个框的最小框的左上角和右下角
enclose_mins = torch.min(b1_mins, b2_mins)
enclose_maxes = torch.max(b1_maxes, b2_maxes)
enclose_wh = torch.max(enclose_maxes - enclose_mins, torch.zeros_like(intersect_maxes))
# 计算对角线距离,最小围框对角线长度
enclose_diagonal = torch.sum(torch.pow(enclose_wh, 2), axis=-1)
ciou = iou - 1.0 * (center_distance) / torch.clamp(enclose_diagonal, min=1e-6)
# 用来测量长宽比的一致性
v = (4 / (math.pi ** 2)) * torch.pow((torch.atan(b1_wh[..., 0] / torch.clamp(b1_wh[..., 1], min=1e-6)) - torch.atan(
b2_wh[..., 0] / torch.clamp(b2_wh[..., 1], min=1e-6))), 2)
alpha = v / torch.clamp((1.0 - iou + v), min=1e-6)
ciou = ciou - alpha * v
return ciou
def BCELoss(pred, target):
"""
交叉熵损失函数
:param pred: shape is (bs,3,in_h,in_w)
:param target: shape is (bs,3,in_h,in_w)
:return: output shape is (bs,3,in_h,in_w)
"""
epsilon = 1e-7
pred = clip_by_tensor(pred, epsilon, 1.0 - epsilon)
output = -target * torch.log(pred) - (1.0 - target) * torch.log(1.0 - pred)
return output
def smooth_labels(y_true, label_smoothing, num_classes):
"""
:param y_true:
:param label_smoothing:
:param num_classes:
:return:
"""
return y_true * (1.0 - label_smoothing) + label_smoothing / num_classes
class YOLOLoss(nn.Module):
def __init__(self, anchors, num_classes, img_size, label_smooth=0, cuda=True):
"""
:param anchors: 锚框，从txt文件中读取的，shape is (9, 2)
:param num_classes: 一共多少分类
:param img_size: 网络输入图片的大小
:param label_smooth:
:param cuda:
"""
super(YOLOLoss, self).__init__()
self.anchors = anchors
self.num_anchors = len(anchors)
self.num_classes = num_classes
self.bbox_attrs = 5 + num_classes
self.img_size = img_size
self.feature_length = [img_size[0] // 32, img_size[0] // 16, img_size[0] // 8]
self.label_smooth = label_smooth
self.ignore_threshold = 0.5
self.lambda_conf = 1.0
self.lambda_cls = 1.0
self.lambda_loc = 1.0
self.cuda = cuda
def forward(self, input, targets=None,device=None):
"""
:param input: yolo网络的输出out0, out1, out2, out0 shape is (N, num_anchors*(5+num_classes), H_in/8, W_in/8)
:param targets: box、label，box中均为归一化的值, type is list
:return: loss: 总损失
loss_conf.item() loss_cls.item(), loss_loc.item()
"""
# 一共多少张图片
bs = input.size(0)
# 特征层的高
in_h = input.size(2)
# 特征层的宽
in_w = input.size(3)
# 计算步长
# 每一个特征点对应原来的图片上多少个像素点
# 如果特征层为13x13的话，一个特征点就对应原来的图片上的32个像素点
stride_h = self.img_size[1] / in_h
stride_w = self.img_size[0] / in_w
# 把先验框的尺寸调整成特征层大小的形式
# 计算出先验框在特征层上对应的宽高
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors]
# bs,3*(5+num_classes),13,13 -> bs,3,13,13,(5+num_classes)
prediction = input.view(bs, int(self.num_anchors / 3),
self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
# 对prediction预测进行调整
conf = torch.sigmoid(prediction[..., 4]) # Conf,是否属于物体的置信度,shape is (bs,3,in_h,in_w)
pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred，框类别的置信度,shape is (bs,3,in_h,in_w,num_classes)
# 找到哪些先验框内部包含物体
# mask shape is (bs,3,in_h,in_w)有物体的点对应位置为1，其余位置为0
# t_box shape is (bs,3,in_h,in_w,4) t_box[...,0] is tx,ty,tw,th，标签框的值，这些值没做归一化
mask, noobj_mask, t_box, tconf, tcls, box_loss_scale_x, box_loss_scale_y = self.get_target(targets,
scaled_anchors,
in_w, in_h,
self.ignore_threshold)
# noobj_mask: shape is (bs,3,in_h,in_w)有物体的点对应位置为0，其余位置为1。是经过进一步处理之后的结果
# pred_boxes_for_ciou shape is (bs, 3, in_h, in_w, 4)，调整以后的预测框中心与宽高
noobj_mask, pred_boxes_for_ciou = self.get_ignore(prediction, targets, scaled_anchors, in_w, in_h, noobj_mask)
if self.cuda: # 将数据放到GPU上
mask, noobj_mask = mask.to(device), noobj_mask.to(device)
box_loss_scale_x, box_loss_scale_y = box_loss_scale_x.to(device), box_loss_scale_y.to(device)
tconf, tcls = tconf.to(device), tcls.to(device)
pred_boxes_for_ciou = pred_boxes_for_ciou.to(device)
t_box = t_box.to(device)
box_loss_scale = 2 - box_loss_scale_x * box_loss_scale_y # shape is (bs, 3, in_h, in_w)
# losses.计算预测框和标签框的CIOU
ciou = (1 - box_ciou(pred_boxes_for_ciou[mask.bool()], t_box[mask.bool()])) * box_loss_scale[mask.bool()]
loss_loc = torch.sum(ciou / bs)
# 计算对应位置框是否有物体的损失,最终得到具体的值
loss_conf = torch.sum(BCELoss(conf, mask) * mask / bs) + torch.sum(BCELoss(conf, mask) * noobj_mask / bs)
# print(smooth_labels(tcls[mask == 1],self.label_smooth,self.num_classes))
# 分类的损失，最终得到具体的值
loss_cls = torch.sum(
BCELoss(pred_cls[mask == 1], smooth_labels(tcls[mask == 1], self.label_smooth, self.num_classes)) / bs)
# print(loss_loc,loss_conf,loss_cls)
loss = loss_conf * self.lambda_conf + loss_cls * self.lambda_cls + loss_loc * self.lambda_loc
return loss, loss_conf.item(), loss_cls.item(), loss_loc.item()
def get_target(self, target, anchors, in_w, in_h, ignore_threshold):
"""
输入图像的box信息、先验框高宽、对应尺度的特征图的高宽，输出一些信息
:param target: box、label，box均为归一化的值,type is list
:param anchors: 先验框在特征层上对应的宽高
:param in_w: yolo输出out特征图的宽 input weight=416时，in_w=(13 or 26 or 52)
:param in_h: yolo输出out特征图的高 input height=416时，in_h=(13 or 26 or 52)
:param ignore_threshold:
:return: mask: shape is (bs,3,in_h,in_w)有物体的点对应位置为1，其余位置为0
noobj_mask: shape is (bs,3,in_h,in_w)有物体的点对应位置为0，其余位置为1
t_box: shape is (bs,3,in_h,in_w,4) t_box[...,0] is tx,ty,tw,th，标签框的值，这些值没做归一化
tconf: shape is (bs,3,in_h,in_w)有物体的点对应位置为1，其余位置为0
tcls: shape is (bs,3,in_h,in_w,num_classes)对应类别的点为1，其余位置为0，相当于做哑编码
box_loss_scale_x: shape is (bs,3,in_h,in_w)，对应网格点中心坐标x的值，此值在0~1之间，也就是标签框里的信息
box_loss_scale_y: shape is (bs,3,in_h,in_w)，对应网格点中心坐标y的值，此值在0~1之间
"""
# 计算一共有多少张图片
bs = len(target) # 即使bs=0，也不影响后续的计算
# 获得先验框
anchor_index = [[0, 1, 2], [3, 4, 5], [6, 7, 8]][self.feature_length.index(in_w)]
subtract_index = [0, 3, 6][self.feature_length.index(in_w)] # 为了找索引，3个一组
# 创建全是0或者全是1的阵列
mask = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
noobj_mask = torch.ones(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
tx = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False) # shape is (N,3,h,w)
ty = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
tw = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
th = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
t_box = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, 4, requires_grad=False)
tconf = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
tcls = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, self.num_classes, requires_grad=False)
box_loss_scale_x = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
box_loss_scale_y = torch.zeros(bs, int(self.num_anchors / 3), in_h, in_w, requires_grad=False)
for b in range(bs): # 对batch_size张图进行遍历，b表示对应图的索引
for t in range(target[b].shape[0]): # 对每张图中的框进行遍历，t来表示图中box的索引
# 计算出在特征层上的点位
gx = target[b][t, 0] * in_w # 中心坐标x轴信息
gy = target[b][t, 1] * in_h
gw = target[b][t, 2] * in_w
gh = target[b][t, 3] * in_h
# 计算出属于哪个网格
gi = int(gx) # int向下取整
gj = int(gy)
# 计算真实框的位置
gt_box = torch.FloatTensor(np.array([0, 0, gw, gh])).unsqueeze(0) # shape is (1, 4)
# 计算出所有先验框的位置
anchor_shapes = torch.FloatTensor(
np.concatenate((np.zeros((self.num_anchors, 2)), np.array(anchors)), 1)) # shape is (9, 4)
# 计算和所有先验框的iou
anch_ious = bbox_iou(gt_box, anchor_shapes) # shape is 9
# Find the best matching anchor box
best_n = np.argmax(anch_ious)
if best_n not in anchor_index: # 保证在对应尺寸的anchor_index中
continue
# Masks
if (gj < in_h) and (gi < in_w): # 确保网格索引没有越界
best_n = best_n - subtract_index #
# 判定哪些先验框内部真实的存在物体
noobj_mask[b, best_n, gj, gi] = 0 # 初始化时，全为1
mask[b, best_n, gj, gi] = 1 # 初始化时，全为0
# 计算标签框中心调整参数
tx[b, best_n, gj, gi] = gx # 对应网格点放置gx，每个框的中心位置不可能相同，所以可以这么操作
ty[b, best_n, gj, gi] = gy
# 计算标签框宽高调整参数
tw[b, best_n, gj, gi] = gw
th[b, best_n, gj, gi] = gh
# 用于获得xywh的比例
box_loss_scale_x[b, best_n, gj, gi] = target[b][t, 2] # 中心坐标x的值，此值在0~1之间
box_loss_scale_y[b, best_n, gj, gi] = target[b][t, 3] # 中心坐标y的值，此值在0~1之间
# 物体置信度
tconf[b, best_n, gj, gi] = 1 # 对应网格位置的置信度为1，说明此处是有物体的
# 种类
tcls[b, best_n, gj, gi, int(target[b][t, 4])] = 1 # 对应维度上的值全为1
else:
print('Step {0} out of bound'.format(b))
print('gj: {0}, height: {1} | gi: {2}, width: {3}'.format(gj, in_h, gi, in_w))
continue
t_box[..., 0] = tx
t_box[..., 1] = ty
t_box[..., 2] = tw
t_box[..., 3] = th
return mask, noobj_mask, t_box, tconf, tcls, box_loss_scale_x, box_loss_scale_y
def get_ignore(self, prediction, target, scaled_anchors, in_w, in_h, noobj_mask):
"""
对预测框的坐标和高宽做进一步转换，通过预测进一步筛选无物体的mask
:param prediction: shape is [bs,3,in_h,in_w,(5+num_classes)]
:param target: box、label，box均为归一化的值,type is list,len(target)=bs
:param scaled_anchors: 先验框在特征层上对应的宽高
:param in_w: yolo输出out特征图的宽 input weight=416时，in_w=(13 or 26 or 52)
:param in_h: yolo输出out特征图的宽 input height=416时，in_w=(13 or 26 or 52)
:param noobj_mask: shape is (bs,3,in_h,in_w)有物体的点对应位置为0，其余位置为1
:return: noobj_mask: shape is (bs,3,in_h,in_w)有物体的点对应位置为0，其余位置为1。是经过进一步处理之后的结果
pred_boxes: shape is (bs, 3, in_h, in_w, 4)，调整以后的预测框中心与宽高
"""
bs = len(target)
anchor_index = [[0, 1, 2], [3, 4, 5], [6, 7, 8]][self.feature_length.index(in_w)]
scaled_anchors = np.array(scaled_anchors)[anchor_index]
# 预测框的中心位置的调整参数
x = torch.sigmoid(prediction[..., 0])
y = torch.sigmoid(prediction[..., 1])
# 预测框的宽高调整参数
w = prediction[..., 2] # Width
h = prediction[..., 3] # Height
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
# 生成网格，先验框中心，网格左上角。t()进行转置
grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_w, 1).repeat(
int(bs * self.num_anchors / 3), 1, 1).view(x.shape).type(FloatTensor)
# grid_x shape is (bs, 3, in_w, in_w)，grid_y shape is (bs, 3, in_h, in_h)
grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_h, 1).t().repeat(
int(bs * self.num_anchors / 3), 1, 1).view(y.shape).type(FloatTensor)
# 生成先验框的宽高，第一个参数是索引的对象，第二个参数0表示按行索引，1表示按列进行索引，第三个参数是一个tensor，就是索引的序号
anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0])) # shape is (3, 1)
anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape) # shape is (bs, 3, in_h, in_w)
anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape) # shape is (bs, 3, in_h, in_w)
# 计算调整预测框中心与宽高
pred_boxes = FloatTensor(prediction[..., :4].shape) # 随机产生值，shape is (bs, 3, in_h, in_w, 4)
pred_boxes[..., 0] = x + grid_x
pred_boxes[..., 1] = y + grid_y
pred_boxes[..., 2] = torch.exp(w) * anchor_w
pred_boxes[..., 3] = torch.exp(h) * anchor_h
for i in range(bs):
pred_boxes_for_ignore = pred_boxes[i]
pred_boxes_for_ignore = pred_boxes_for_ignore.view(-1, 4)
if len(target[i]) > 0:
gx = target[i][:, 0:1] * in_w
gy = target[i][:, 1:2] * in_h
gw = target[i][:, 2:3] * in_w
gh = target[i][:, 3:4] * in_h
gt_box = torch.FloatTensor(np.concatenate([gx, gy, gw, gh], -1)).type(FloatTensor) # shape is (N, 4)
# inter/union shape is (N, M), N is number of gt_box, M is number of pred_boxes_for_ignore,M=3*in_h*in_w
# 计算一张图中所有标签框和预测框pred_boxes_for_ignore的IOU
anch_ious = jaccard(gt_box, pred_boxes_for_ignore)
for t in range(target[i].shape[0]): # t表示每个图中的标签box的索引，最大值为N-1
anch_iou = anch_ious[t].view(pred_boxes[i].size()[:3]) # shape is (3, in_h, in_w)
noobj_mask[i][anch_iou > self.ignore_threshold] = 0
return noobj_mask, pred_boxes
train
from model import YoloBody
from config import parser
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from torch.utils.data import DataLoader
from datalist import YoloDataset,yolo_dataset_collate
import numpy as np
from utils import YOLOLoss
the_loss=1000
class train(object):
def __init__(self):
self.args=parser.parse_args()
print(f"-----------{self.args.project_name}-------------")
use_cuda=self.args.use_cuda and torch.cuda.is_available()
if use_cuda:
torch.cuda.manual_seed(self.args.seed)
torch.cuda.manual_seed_all(self.args.seed)
else:
torch.manual_seed(self.args.seed)
self.device=torch.device('cuda' if use_cuda else 'cpu')
train_kwargs={'num_workers':0,'pin_memory':True} if use_cuda else {}
test_kwargs={'num_workers':0,'pin_memory':False} if use_cuda else {}
'''
构造DataLoader
'''
self.annotations=self.args.data_path
with open(self.annotations) as f:
self.lines=f.readlines()
np.random.seed(self.args.seed)
np.random.shuffle(self.lines)
np.random.seed(None)
self.num_val=int(len(self.lines)*self.args.val_num)
self.num_train=len(self.lines)-self.num_val
self.train_dataset=YoloDataset(self.lines[:self.num_train],self.args.image_size)
self.test_dataset=YoloDataset(self.lines[self.num_train:],self.args.image_size)
self.train_dataloader=DataLoader(self.train_dataset,batch_size=self.args.train_batchsize,shuffle=True,**train_kwargs,collate_fn=yolo_dataset_collate)
self.test_dataloader=DataLoader(self.test_dataset,batch_size=self.args.test_batchsize,shuffle=False,**test_kwargs,collate_fn=yolo_dataset_collate)
'''
构造model
'''
self.num_classes=self.get_classes(self.args.classes_path)
self.anchors=self.get_anchors(self.args.anchor_path)
self.model=YoloBody(len(self.anchors[0]),len(self.num_classes))
if use_cuda:
self.model=torch.nn.DataParallel(self.model,device_ids=range(torch.cuda.device_count()))
cudnn.benchmark=True
if self.args.resume:
try:
print("load the weight from pretrained-weight file")
model_dict = self.model.state_dict()
pretrained_dict = torch.load(self.args.pretrained_weight, map_location=self.device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
self.model.load_state_dict(model_dict)
print("Finished to load the weight")
except:
print("can not load weight \n train the model from stratch")
self.model.apply(self.weights_init)
'''
构造loss目标函数
选择优化器
学习率变化选择
'''
self.criterion=self.create_loss()
self.optimizer=optim.Adam(self.model.parameters(),self.args.lr)
self.scheduler=optim.lr_scheduler.CosineAnnealingLR(self.optimizer,T_max=5,eta_min=1e-5)
for epoch in range(1,self.args.epoches):
self.train(epoch)
if epoch %1==0:
self.test(epoch)
torch.cuda.empty_cache()
print("finish model training")
def train(self,epoch):
self.model.train()
losses = []
pbar=tqdm(self.train_dataloader,desc=f'Train Epoch{epoch}/{self.args.epoches}')
for data, target in pbar:
total_loss = 0
data,target=torch.from_numpy(data).type(torch.FloatTensor), [torch.from_numpy(ann).type(torch.FloatTensor) for ann in target]
data,target=data.to(self.device),target
self.optimizer.zero_grad()
outputs=self.model(data)
for i in range(3):
loss_item=self.criterion[i](outputs[i],target,self.device)[0]
total_loss+=loss_item
losses.append(loss_item.item())
loss=sum(losses)
total_loss.backward()
self.optimizer.step()
pbar.set_description(
f'Train Epoch:{epoch}/{self.args.epoches} train_loss:{round(np.mean(losses),4)}'
)
self.scheduler.step()
def test(self,epoch):
self.model.eval()
losses=[]
pbar=tqdm(self.test_dataloader,desc=f'Test Epoch{epoch}/{self.args.epoches}')
with torch.no_grad():
for data,target in pbar:
data, target = torch.from_numpy(data).type(torch.FloatTensor), [
torch.from_numpy(ann).type(torch.FloatTensor) for ann in target]
data,target=data.to(self.device),target
outputs=self.model(data)
for i in range(3):
loss_item=self.criterion[i](outputs[i],target)[0]
losses.append(loss_item.item())
pbar.set_description(
f'【Test】 Epoch:{epoch}/{self.args.epoches} test_loss:{round(np.mean(losses),4)}'
)
global the_loss
if np.mean(losses)>the_loss:
the_loss=np.mean(losses)
torch.save({
'epoch': epoch,
'model_state_dict': self.model.state_dict(),
"optimizer_state_dict": self.optimizer.state_dict(),
}
,'./weights/'+f'Epoch_'+str(epoch)+'_Loss_'+str(the_loss)+'.pth'
)
def get_classes(self,classes_path):
with open(classes_path) as f:
class_names=f.readlines()
class_names=[c.strip() for c in class_names]
return class_names
def get_anchors(self,anchors_path):
with open(anchors_path) as f:
anchors=f.readline()
anchors=[float(x) for x in anchors.split(',')]
return np.array(anchors).reshape(-1,3,2)[::-1,:,:]
#因为这里的loss比较复杂
def create_loss(self):
yolo_losses=[]
for i in range(3):
yolo_losses.append(YOLOLoss(np.reshape(self.anchors,[-1,2]),len(self.num_classes),
(416,416)))
return yolo_losses
train=train()
总结
yolo系列是一个作为cver的工程师很好的入门方向，尤其是yolov4中集成了大量的trick。但我在这里把能省去的都省去了，学习的话，可以在此基础之上增加相关内容，我这边也是尽可能把代码结构简化，风格严格控制，我想这是有利于学习的，因为我感觉每次去github抄代码，一个人就是一个哈姆雷特，学起来太难了，作为一个算法工程师，再不济也要坚持自己撸代码
3万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
