updates
This commit is contained in:
parent
d2ef817b1f
commit
5f2b551818
28
models.py
28
models.py
|
@ -4,13 +4,11 @@ from utils.parse_config import *
|
||||||
from utils.utils import *
|
from utils.utils import *
|
||||||
|
|
||||||
ONNX_EXPORT = False
|
ONNX_EXPORT = False
|
||||||
arc = 'normal' # (normal, uCE, uBCE, uBCEs) detection architectures
|
|
||||||
|
|
||||||
|
|
||||||
def create_modules(module_defs, img_size):
|
def create_modules(module_defs, img_size, arc):
|
||||||
"""
|
# Constructs module list of layer blocks from module configuration in module_defs
|
||||||
Constructs module list of layer blocks from module configuration in module_defs
|
|
||||||
"""
|
|
||||||
hyperparams = module_defs.pop(0)
|
hyperparams = module_defs.pop(0)
|
||||||
output_filters = [int(hyperparams['channels'])]
|
output_filters = [int(hyperparams['channels'])]
|
||||||
module_list = nn.ModuleList()
|
module_list = nn.ModuleList()
|
||||||
|
@ -74,11 +72,12 @@ def create_modules(module_defs, img_size):
|
||||||
modules = YOLOLayer(anchors=mdef['anchors'][mask], # anchor list
|
modules = YOLOLayer(anchors=mdef['anchors'][mask], # anchor list
|
||||||
nc=int(mdef['classes']), # number of classes
|
nc=int(mdef['classes']), # number of classes
|
||||||
img_size=img_size, # (416, 416)
|
img_size=img_size, # (416, 416)
|
||||||
yolo_index=yolo_index) # 0, 1 or 2
|
yolo_index=yolo_index, # 0, 1 or 2
|
||||||
|
arc=arc) # yolo architecture
|
||||||
|
|
||||||
# Initialize preceding Conv2d() bias (https://arxiv.org/pdf/1708.02002.pdf section 3.3)
|
# Initialize preceding Conv2d() bias (https://arxiv.org/pdf/1708.02002.pdf section 3.3)
|
||||||
try:
|
try:
|
||||||
if arc == 'normal':
|
if arc == 'default':
|
||||||
b = [-5.0, -4.0] # obj, cls
|
b = [-5.0, -4.0] # obj, cls
|
||||||
elif arc == 'uCE': # unified CE (1 background + 80 classes)
|
elif arc == 'uCE': # unified CE (1 background + 80 classes)
|
||||||
b = [3.0, -4.0] # obj, cls
|
b = [3.0, -4.0] # obj, cls
|
||||||
|
@ -113,7 +112,7 @@ class Swish(nn.Module):
|
||||||
|
|
||||||
|
|
||||||
class YOLOLayer(nn.Module):
|
class YOLOLayer(nn.Module):
|
||||||
def __init__(self, anchors, nc, img_size, yolo_index):
|
def __init__(self, anchors, nc, img_size, yolo_index, arc):
|
||||||
super(YOLOLayer, self).__init__()
|
super(YOLOLayer, self).__init__()
|
||||||
|
|
||||||
self.anchors = torch.Tensor(anchors)
|
self.anchors = torch.Tensor(anchors)
|
||||||
|
@ -121,6 +120,7 @@ class YOLOLayer(nn.Module):
|
||||||
self.nc = nc # number of classes (80)
|
self.nc = nc # number of classes (80)
|
||||||
self.nx = 0 # initialize number of x gridpoints
|
self.nx = 0 # initialize number of x gridpoints
|
||||||
self.ny = 0 # initialize number of y gridpoints
|
self.ny = 0 # initialize number of y gridpoints
|
||||||
|
self.arc = arc
|
||||||
|
|
||||||
if ONNX_EXPORT: # grids must be computed in __init__
|
if ONNX_EXPORT: # grids must be computed in __init__
|
||||||
stride = [32, 16, 8][yolo_index] # stride of this layer
|
stride = [32, 16, 8][yolo_index] # stride of this layer
|
||||||
|
@ -175,12 +175,12 @@ class YOLOLayer(nn.Module):
|
||||||
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
|
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
|
||||||
io[..., :4] *= self.stride
|
io[..., :4] *= self.stride
|
||||||
|
|
||||||
if arc == 'normal':
|
if self.arc == 'default':
|
||||||
torch.sigmoid_(io[..., 4:])
|
torch.sigmoid_(io[..., 4:])
|
||||||
elif arc == 'uCE': # unified CE (1 background + 80 classes)
|
elif self.arc == 'uCE': # unified CE (1 background + 80 classes)
|
||||||
io[..., 4:] = F.softmax(io[..., 4:], dim=4)
|
io[..., 4:] = F.softmax(io[..., 4:], dim=4)
|
||||||
io[..., 4] = 1
|
io[..., 4] = 1
|
||||||
elif arc == 'uBCE': # unified BCE (80 classes)
|
elif self.arc == 'uBCE': # unified BCE (80 classes)
|
||||||
torch.sigmoid_(io[..., 5:])
|
torch.sigmoid_(io[..., 5:])
|
||||||
io[..., 4] = 1
|
io[..., 4] = 1
|
||||||
|
|
||||||
|
@ -192,13 +192,13 @@ class YOLOLayer(nn.Module):
|
||||||
|
|
||||||
|
|
||||||
class Darknet(nn.Module):
|
class Darknet(nn.Module):
|
||||||
"""YOLOv3 object detection model"""
|
# YOLOv3 object detection model
|
||||||
|
|
||||||
def __init__(self, cfg, img_size=(416, 416)):
|
def __init__(self, cfg, img_size=(416, 416), arc='default'):
|
||||||
super(Darknet, self).__init__()
|
super(Darknet, self).__init__()
|
||||||
|
|
||||||
self.module_defs = parse_model_cfg(cfg)
|
self.module_defs = parse_model_cfg(cfg)
|
||||||
self.module_list, self.routs = create_modules(self.module_defs, img_size)
|
self.module_list, self.routs = create_modules(self.module_defs, img_size, arc)
|
||||||
self.yolo_layers = get_yolo_layers(self)
|
self.yolo_layers = get_yolo_layers(self)
|
||||||
|
|
||||||
# Darknet Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
|
# Darknet Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
|
||||||
|
|
5
train.py
5
train.py
|
@ -84,7 +84,7 @@ def train():
|
||||||
nc = int(data_dict['classes']) # number of classes
|
nc = int(data_dict['classes']) # number of classes
|
||||||
|
|
||||||
# Initialize model
|
# Initialize model
|
||||||
model = Darknet(cfg).to(device)
|
model = Darknet(cfg, arc=opt.arc).to(device)
|
||||||
|
|
||||||
# Optimizer
|
# Optimizer
|
||||||
# optimizer = optim.Adam(model.parameters(), lr=hyp['lr0'], weight_decay=hyp['weight_decay'])
|
# optimizer = optim.Adam(model.parameters(), lr=hyp['lr0'], weight_decay=hyp['weight_decay'])
|
||||||
|
@ -259,7 +259,7 @@ def train():
|
||||||
pred = model(imgs)
|
pred = model(imgs)
|
||||||
|
|
||||||
# Compute loss
|
# Compute loss
|
||||||
loss, loss_items = compute_loss(pred, targets, model)
|
loss, loss_items = compute_loss(pred, targets, model, arc=opt.arc)
|
||||||
if torch.isnan(loss):
|
if torch.isnan(loss):
|
||||||
print('WARNING: nan loss detected, ending training')
|
print('WARNING: nan loss detected, ending training')
|
||||||
return results
|
return results
|
||||||
|
@ -367,6 +367,7 @@ if __name__ == '__main__':
|
||||||
parser.add_argument('--img-weights', action='store_true', help='select training images by weight')
|
parser.add_argument('--img-weights', action='store_true', help='select training images by weight')
|
||||||
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
|
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
|
||||||
parser.add_argument('--weights', type=str, default='', help='initial weights') # i.e. weights/darknet.53.conv.74
|
parser.add_argument('--weights', type=str, default='', help='initial weights') # i.e. weights/darknet.53.conv.74
|
||||||
|
parser.add_argument('--arc', type=str, default='default', help='yolo architecture') # default, uCE, uBCE
|
||||||
opt = parser.parse_args()
|
opt = parser.parse_args()
|
||||||
opt.weights = 'weights/last.pt' if opt.resume else opt.weights
|
opt.weights = 'weights/last.pt' if opt.resume else opt.weights
|
||||||
print(opt)
|
print(opt)
|
||||||
|
|
|
@ -312,7 +312,7 @@ class FocalLoss(nn.Module):
|
||||||
return loss
|
return loss
|
||||||
|
|
||||||
|
|
||||||
def compute_loss(p, targets, model): # predictions, targets, model
|
def compute_loss(p, targets, model, arc='default'): # predictions, targets, model
|
||||||
ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor
|
ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor
|
||||||
lcls, lbox, lobj = ft([0]), ft([0]), ft([0])
|
lcls, lbox, lobj = ft([0]), ft([0]), ft([0])
|
||||||
tcls, tbox, indices, anchor_vec = build_targets(model, targets)
|
tcls, tbox, indices, anchor_vec = build_targets(model, targets)
|
||||||
|
@ -321,12 +321,12 @@ def compute_loss(p, targets, model): # predictions, targets, model
|
||||||
# Define criteria
|
# Define criteria
|
||||||
BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]))
|
BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]))
|
||||||
BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]))
|
BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]))
|
||||||
# CE = nn.CrossEntropyLoss(weight=model.class_weights)
|
BCE = nn.BCEWithLogitsLoss()
|
||||||
|
CE = nn.CrossEntropyLoss() # weight=model.class_weights
|
||||||
|
|
||||||
# Compute losses
|
# Compute losses
|
||||||
bs = p[0].shape[0] # batch size
|
bs = p[0].shape[0] # batch size
|
||||||
k = bs / 64 # loss gain
|
k = bs / 64 # loss gain
|
||||||
arc = 'normal' # (normal, uCE, uBCE, uBCEs) detection architectures
|
|
||||||
for i, pi in enumerate(p): # layer index, layer predictions
|
for i, pi in enumerate(p): # layer index, layer predictions
|
||||||
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
||||||
tobj = torch.zeros_like(pi[..., 0]) # target obj
|
tobj = torch.zeros_like(pi[..., 0]) # target obj
|
||||||
|
@ -344,7 +344,7 @@ def compute_loss(p, targets, model): # predictions, targets, model
|
||||||
giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation
|
giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation
|
||||||
lbox += (1.0 - giou).mean() # giou loss
|
lbox += (1.0 - giou).mean() # giou loss
|
||||||
|
|
||||||
if arc == 'normal' and model.nc > 1: # cls loss (only if multiple classes)
|
if arc == 'default' and model.nc > 1: # cls loss (only if multiple classes)
|
||||||
t = torch.zeros_like(ps[:, 5:]) # targets
|
t = torch.zeros_like(ps[:, 5:]) # targets
|
||||||
t[range(nb), tcls[i]] = 1.0
|
t[range(nb), tcls[i]] = 1.0
|
||||||
lcls += BCEcls(ps[:, 5:], t) # BCE
|
lcls += BCEcls(ps[:, 5:], t) # BCE
|
||||||
|
@ -354,7 +354,7 @@ def compute_loss(p, targets, model): # predictions, targets, model
|
||||||
# with open('targets.txt', 'a') as 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)]
|
# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
|
||||||
|
|
||||||
if arc == 'normal':
|
if arc == 'default': # (default, uCE, uBCE) detection architectures
|
||||||
lobj += BCEobj(pi[..., 4], tobj) # obj loss
|
lobj += BCEobj(pi[..., 4], tobj) # obj loss
|
||||||
|
|
||||||
elif arc == 'uCE': # unified CE (1 background + 80 classes), hyps 20
|
elif arc == 'uCE': # unified CE (1 background + 80 classes), hyps 20
|
||||||
|
@ -367,7 +367,7 @@ def compute_loss(p, targets, model): # predictions, targets, model
|
||||||
t = torch.zeros_like(pi[..., 5:]) # targets
|
t = torch.zeros_like(pi[..., 5:]) # targets
|
||||||
if nb:
|
if nb:
|
||||||
t[b, a, gj, gi, tcls[i]] = 1.0
|
t[b, a, gj, gi, tcls[i]] = 1.0
|
||||||
lobj += BCEobj(pi[..., 5:], t)
|
lobj += BCE(pi[..., 5:], t)
|
||||||
|
|
||||||
lbox *= k * h['giou']
|
lbox *= k * h['giou']
|
||||||
lobj *= k * h['obj']
|
lobj *= k * h['obj']
|
||||||
|
|
Loading…
Reference in New Issue