car-detection-bayes/models.py

470 lines
21 KiB
Python
Raw Normal View History

2019-09-16 21:15:07 +00:00
from utils.google_utils import *
2020-04-02 19:22:15 +00:00
from utils.layers import *
2018-08-26 08:51:39 +00:00
from utils.parse_config import *
2019-07-29 10:06:29 +00:00
2019-08-11 13:22:53 +00:00
ONNX_EXPORT = False
2019-01-03 22:41:31 +00:00
2018-08-26 08:51:39 +00:00
def create_modules(module_defs, img_size):
2019-08-23 15:18:59 +00:00
# Constructs module list of layer blocks from module configuration in module_defs
2020-04-03 02:10:51 +00:00
img_size = [img_size] * 2 if isinstance(img_size, int) else img_size # expand if necessary
_ = module_defs.pop(0) # cfg training hyperparams (unused)
output_filters = [3] # input channels
2018-08-26 08:51:39 +00:00
module_list = nn.ModuleList()
2019-12-29 22:54:08 +00:00
routs = [] # list of layers which rout to deeper layers
2019-07-03 12:42:11 +00:00
yolo_index = -1
2019-08-12 11:37:11 +00:00
for i, mdef in enumerate(module_defs):
2018-08-26 08:51:39 +00:00
modules = nn.Sequential()
2019-08-03 12:49:38 +00:00
if mdef['type'] == 'convolutional':
2020-02-20 01:08:03 +00:00
bn = mdef['batch_normalize']
filters = mdef['filters']
2020-04-11 19:37:03 +00:00
k = mdef['size'] # kernel size
2020-02-20 01:08:03 +00:00
stride = mdef['stride'] if 'stride' in mdef else (mdef['stride_y'], mdef['stride_x'])
2020-04-11 19:37:03 +00:00
if isinstance(k, int): # single-size conv
2020-04-04 03:03:44 +00:00
modules.add_module('Conv2d', nn.Conv2d(in_channels=output_filters[-1],
out_channels=filters,
2020-04-11 19:37:03 +00:00
kernel_size=k,
2020-04-04 03:03:44 +00:00
stride=stride,
2020-04-11 19:37:03 +00:00
padding=k // 2 if mdef['pad'] else 0,
2020-04-04 03:03:44 +00:00
groups=mdef['groups'] if 'groups' in mdef else 1,
bias=not bn))
else: # multiple-size conv
modules.add_module('MixConv2d', MixConv2d(in_ch=output_filters[-1],
out_ch=filters,
2020-04-11 19:37:03 +00:00
k=k,
2020-04-04 03:03:44 +00:00
stride=stride,
bias=not bn))
2018-08-26 08:51:39 +00:00
if bn:
2020-04-02 21:10:45 +00:00
modules.add_module('BatchNorm2d', nn.BatchNorm2d(filters, momentum=0.03, eps=1E-4))
2020-03-10 01:49:35 +00:00
else:
routs.append(i) # detection output (goes into yolo layer)
2020-02-28 18:06:35 +00:00
if mdef['activation'] == 'leaky': # activation study https://github.com/ultralytics/yolov3/issues/441
2019-08-03 12:38:06 +00:00
modules.add_module('activation', nn.LeakyReLU(0.1, inplace=True))
2019-08-11 13:17:40 +00:00
# modules.add_module('activation', nn.PReLU(num_parameters=1, init=0.10))
2019-11-23 02:20:11 +00:00
elif mdef['activation'] == 'swish':
modules.add_module('activation', Swish())
2018-08-26 08:51:39 +00:00
elif mdef['type'] == 'BatchNorm2d':
filters = output_filters[-1]
modules = nn.BatchNorm2d(filters, momentum=0.03, eps=1E-4)
if i == 0 and filters == 3: # normalize RGB image
# imagenet mean and var https://pytorch.org/docs/stable/torchvision/models.html#classification
modules.running_mean = torch.tensor([0.485, 0.456, 0.406])
modules.running_var = torch.tensor([0.0524, 0.0502, 0.0506])
2019-08-03 12:49:38 +00:00
elif mdef['type'] == 'maxpool':
2020-04-11 19:37:03 +00:00
k = mdef['size'] # kernel size
2020-02-20 01:08:03 +00:00
stride = mdef['stride']
2020-04-11 19:37:03 +00:00
maxpool = nn.MaxPool2d(kernel_size=k, stride=stride, padding=(k - 1) // 2)
if k == 2 and stride == 1: # yolov3-tiny
2019-08-03 12:38:06 +00:00
modules.add_module('ZeroPad2d', nn.ZeroPad2d((0, 1, 0, 1)))
modules.add_module('MaxPool2d', maxpool)
else:
modules = maxpool
2018-12-22 11:36:33 +00:00
2019-08-03 12:49:38 +00:00
elif mdef['type'] == 'upsample':
2020-01-30 20:39:54 +00:00
if ONNX_EXPORT: # explicitly state size, avoid scale_factor
2020-02-17 20:36:11 +00:00
g = (yolo_index + 1) * 2 / 32 # gain
modules = nn.Upsample(size=tuple(int(x * g) for x in img_size)) # img_size = (320, 192)
2020-01-30 20:39:54 +00:00
else:
2020-02-20 01:08:03 +00:00
modules = nn.Upsample(scale_factor=mdef['stride'])
2018-08-26 08:51:39 +00:00
2019-08-03 12:49:38 +00:00
elif mdef['type'] == 'route': # nn.Sequential() placeholder for 'route' layer
2020-02-20 01:08:03 +00:00
layers = mdef['layers']
2020-03-19 23:23:44 +00:00
filters = sum([output_filters[l + 1 if l > 0 else l] for l in layers])
routs.extend([i + l if l < 0 else l for l in layers])
2020-04-05 21:47:41 +00:00
modules = FeatureConcat(layers=layers)
2019-08-03 12:14:10 +00:00
2019-08-03 12:49:38 +00:00
elif mdef['type'] == 'shortcut': # nn.Sequential() placeholder for 'shortcut' layer
2020-02-20 01:08:03 +00:00
layers = mdef['from']
2020-02-20 02:26:45 +00:00
filters = output_filters[-1]
2020-02-17 23:28:11 +00:00
routs.extend([i + l if l < 0 else l for l in layers])
2020-04-02 19:22:15 +00:00
modules = WeightedFeatureFusion(layers=layers, weight='weights_type' in mdef)
2018-08-26 08:51:39 +00:00
2019-08-03 12:49:38 +00:00
elif mdef['type'] == 'reorg3d': # yolov3-spp-pan-scale
2019-07-29 10:06:29 +00:00
pass
2019-08-03 12:49:38 +00:00
elif mdef['type'] == 'yolo':
2019-07-03 12:42:11 +00:00
yolo_index += 1
2020-04-03 02:10:51 +00:00
stride = [32, 16, 8, 4, 2][yolo_index] # P3-P7 stride
2020-04-03 19:42:09 +00:00
layers = mdef['from'] if 'from' in mdef else []
2020-03-10 01:24:20 +00:00
modules = YOLOLayer(anchors=mdef['anchors'][mdef['mask']], # anchor list
2020-02-20 01:08:03 +00:00
nc=mdef['classes'], # number of classes
2019-08-12 11:49:38 +00:00
img_size=img_size, # (416, 416)
2020-03-10 01:55:17 +00:00
yolo_index=yolo_index, # 0, 1, 2...
2020-04-03 19:42:09 +00:00
layers=layers, # output layers
2020-04-03 02:10:51 +00:00
stride=stride)
2019-08-18 23:27:41 +00:00
2019-08-19 15:07:16 +00:00
# Initialize preceding Conv2d() bias (https://arxiv.org/pdf/1708.02002.pdf section 3.3)
2019-08-22 21:41:51 +00:00
try:
2020-04-03 19:42:09 +00:00
j = layers[yolo_index] if 'from' in mdef else -1
2020-03-10 03:08:19 +00:00
bias_ = module_list[j][0].bias # shape(255,)
bias = bias_[:modules.no * modules.na].view(modules.na, -1) # shape(3,85)
2020-04-15 18:50:54 +00:00
bias[:, 4] += -4.5 # obj
bias[:, 5:] += math.log(0.6 / (modules.nc - 0.99)) # cls (sigmoid(p) = 1/nc)
2020-03-10 03:08:19 +00:00
module_list[j][0].bias = torch.nn.Parameter(bias_, requires_grad=bias_.requires_grad)
2019-08-22 21:41:51 +00:00
except:
print('WARNING: smart bias initialization failure.')
2019-08-19 15:07:16 +00:00
2019-07-28 22:42:03 +00:00
else:
2019-08-03 12:49:38 +00:00
print('Warning: Unrecognized Layer Type: ' + mdef['type'])
2018-08-26 08:51:39 +00:00
# Register module list and number of output filters
module_list.append(modules)
output_filters.append(filters)
2020-03-19 23:41:42 +00:00
routs_binary = [False] * (i + 1)
for i in routs:
routs_binary[i] = True
return module_list, routs_binary
2018-08-26 08:51:39 +00:00
class YOLOLayer(nn.Module):
2020-04-03 02:10:51 +00:00
def __init__(self, anchors, nc, img_size, yolo_index, layers, stride):
2018-08-26 08:51:39 +00:00
super(YOLOLayer, self).__init__()
2019-08-17 12:15:27 +00:00
self.anchors = torch.Tensor(anchors)
2020-03-10 01:55:17 +00:00
self.index = yolo_index # index of this layer in layers
self.layers = layers # model output layer indices
2020-04-03 02:10:51 +00:00
self.stride = stride # layer stride
2020-03-10 01:55:17 +00:00
self.nl = len(layers) # number of output layers (3)
2019-04-19 18:41:18 +00:00
self.na = len(anchors) # number of anchors (3)
self.nc = nc # number of classes (80)
2020-03-10 01:55:17 +00:00
self.no = nc + 5 # number of outputs (85)
2020-04-12 20:02:00 +00:00
self.nx, self.ny, self.ng = 0, 0, 0 # initialize number of x, y gridpoints
2020-04-03 02:10:51 +00:00
self.anchor_vec = self.anchors / self.stride
self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2)
2018-08-26 08:51:39 +00:00
2020-03-09 17:46:59 +00:00
if ONNX_EXPORT:
2020-04-08 17:25:52 +00:00
self.training = False
2020-04-03 02:10:51 +00:00
self.create_grids((img_size[1] // stride, img_size[0] // stride)) # number x, y grid points
def create_grids(self, ng=(13, 13), device='cpu'):
self.nx, self.ny = ng # x and y grid size
self.ng = torch.tensor(ng)
2020-04-03 02:10:51 +00:00
# build xy offsets
2020-04-03 03:23:55 +00:00
if not self.training:
yv, xv = torch.meshgrid([torch.arange(self.ny, device=device), torch.arange(self.nx, device=device)])
2020-04-03 19:38:08 +00:00
self.grid = torch.stack((xv, yv), 2).view((1, 1, self.ny, self.nx, 2)).float()
2020-04-03 02:10:51 +00:00
if self.anchor_vec.device != device:
self.anchor_vec = self.anchor_vec.to(device)
self.anchor_wh = self.anchor_wh.to(device)
2019-02-19 15:11:18 +00:00
2020-04-14 01:25:59 +00:00
def forward(self, p, out):
2020-03-15 00:04:38 +00:00
ASFF = False # https://arxiv.org/abs/1911.09516
if ASFF:
i, n = self.index, self.nl # index in layers, number of layers
p = out[self.layers[i]]
bs, _, ny, nx = p.shape # bs, 255, 13, 13
if (self.nx, self.ny) != (nx, ny):
2020-04-03 02:10:51 +00:00
self.create_grids((nx, ny), p.device)
2020-03-15 00:04:38 +00:00
# outputs and weights
# w = F.softmax(p[:, -n:], 1) # normalized weights
w = torch.sigmoid(p[:, -n:]) * (2 / n) # sigmoid weights (faster)
# w = w / w.sum(1).unsqueeze(1) # normalize across layer dimension
# weighted ASFF sum
p = out[self.layers[i]][:, :-n] * w[:, i:i + 1]
for j in range(n):
if j != i:
p += w[:, j:j + 1] * \
F.interpolate(out[self.layers[j]][:, :-n], size=[ny, nx], mode='bilinear', align_corners=False)
elif ONNX_EXPORT:
2019-04-21 19:07:01 +00:00
bs = 1 # batch size
2019-02-19 18:00:44 +00:00
else:
2019-12-09 02:08:19 +00:00
bs, _, ny, nx = p.shape # bs, 255, 13, 13
2019-04-25 18:50:37 +00:00
if (self.nx, self.ny) != (nx, ny):
2020-04-03 02:10:51 +00:00
self.create_grids((nx, ny), p.device)
2018-08-26 08:51:39 +00:00
# p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85) # (bs, anchors, grid, grid, classes + xywh)
2019-12-20 02:09:13 +00:00
p = p.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction
2018-08-26 08:51:39 +00:00
if self.training:
return p
elif ONNX_EXPORT:
2020-02-09 05:51:31 +00:00
# Avoid broadcasting for ANE operations
2019-12-04 01:22:58 +00:00
m = self.na * self.nx * self.ny
2020-02-09 19:17:31 +00:00
ng = 1 / self.ng.repeat((m, 1))
2020-04-03 03:23:55 +00:00
grid = self.grid.repeat((1, self.na, 1, 1, 1)).view(m, 2)
2020-02-09 19:17:31 +00:00
anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(m, 2) * ng
2019-12-20 02:09:13 +00:00
p = p.view(m, self.no)
2020-04-03 03:23:55 +00:00
xy = torch.sigmoid(p[:, 0:2]) + grid # x, y
2020-01-11 21:11:30 +00:00
wh = torch.exp(p[:, 2:4]) * anchor_wh # width, height
2020-01-30 05:52:00 +00:00
p_cls = torch.sigmoid(p[:, 4:5]) if self.nc == 1 else \
torch.sigmoid(p[:, 5:self.no]) * torch.sigmoid(p[:, 4:5]) # conf
2020-02-09 19:17:31 +00:00
return p_cls, xy * ng, wh
else: # inference
2019-04-05 13:34:42 +00:00
io = p.clone() # inference output
2020-04-03 03:23:55 +00:00
io[..., :2] = torch.sigmoid(io[..., :2]) + self.grid # xy
2019-04-05 13:34:42 +00:00
io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh # wh yolo method
io[..., :4] *= self.stride
2020-03-09 23:44:26 +00:00
torch.sigmoid_(io[..., 4:])
return io.view(bs, -1, self.no), p # view [1, 3, 13, 13, 85] as [1, 507, 85]
2018-08-26 08:51:39 +00:00
class Darknet(nn.Module):
2019-08-23 15:18:59 +00:00
# YOLOv3 object detection model
2018-08-26 08:51:39 +00:00
def __init__(self, cfg, img_size=(416, 416), verbose=False):
2018-08-26 08:51:39 +00:00
super(Darknet, self).__init__()
2018-12-15 20:06:39 +00:00
2019-04-23 14:48:47 +00:00
self.module_defs = parse_model_cfg(cfg)
self.module_list, self.routs = create_modules(self.module_defs, img_size)
2020-04-14 00:56:12 +00:00
self.yolo_layers = get_yolo_layers(self)
2020-04-11 17:45:33 +00:00
# torch_utils.initialize_weights(self)
2018-08-26 08:51:39 +00:00
2019-06-05 11:49:56 +00:00
# Darknet Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
self.version = np.array([0, 2, 5], dtype=np.int32) # (int32) version info: major, minor, revision
self.seen = np.array([0], dtype=np.int64) # (int64) number of images seen during training
2020-04-15 19:12:59 +00:00
self.info(verbose) if not ONNX_EXPORT else None # print model description
2019-04-23 14:48:47 +00:00
def forward(self, x, augment=False, verbose=False):
2020-04-07 21:19:43 +00:00
if not augment:
return self.forward_once(x)
else: # Augment images (inference and test only) https://github.com/ultralytics/yolov3/issues/931
img_size = x.shape[-2:] # height, width
2020-04-07 21:23:31 +00:00
s = [0.83, 0.67] # scales
2020-04-07 21:19:43 +00:00
y = []
for i, xi in enumerate((x,
torch_utils.scale_img(x.flip(3), s[0], same_shape=False), # flip-lr and scale
torch_utils.scale_img(x, s[1], same_shape=False), # scale
)):
2020-04-07 21:23:31 +00:00
# cv2.imwrite('img%g.jpg' % i, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1])
2020-04-07 21:19:43 +00:00
y.append(self.forward_once(xi)[0])
y[1][..., :4] /= s[0] # scale
y[1][..., 0] = img_size[1] - y[1][..., 0] # flip lr
y[2][..., :4] /= s[1] # scale
# for i, yi in enumerate(y): # coco small, medium, large = < 32**2 < 96**2 <
# area = yi[..., 2:4].prod(2)[:, :, None]
# if i == 1:
2020-04-08 00:35:35 +00:00
# yi *= (area < 96. ** 2).float()
2020-04-07 21:19:43 +00:00
# elif i == 2:
2020-04-08 00:35:35 +00:00
# yi *= (area > 32. ** 2).float()
2020-04-07 21:19:43 +00:00
# y[i] = yi
y = torch.cat(y, 1)
return y, None
def forward_once(self, x, augment=False, verbose=False):
img_size = x.shape[-2:] # height, width
2020-02-19 23:16:00 +00:00
yolo_out, out = [], []
2020-01-23 21:52:17 +00:00
if verbose:
print('0', x.shape)
2020-04-05 22:22:32 +00:00
str = ''
2018-08-26 08:51:39 +00:00
# Augment images (inference and test only)
if augment: # https://github.com/ultralytics/yolov3/issues/931
nb = x.shape[0] # batch size
2020-04-07 21:19:43 +00:00
s = [0.83, 0.67] # scales
x = torch.cat((x,
2020-04-07 21:19:43 +00:00
torch_utils.scale_img(x.flip(3), s[0]), # flip-lr and scale
torch_utils.scale_img(x, s[1]), # scale
), 0)
2020-04-05 22:22:32 +00:00
for i, module in enumerate(self.module_list):
name = module.__class__.__name__
if name in ['WeightedFeatureFusion', 'FeatureConcat']: # sum, concat
2020-02-19 20:59:56 +00:00
if verbose:
2020-02-20 02:26:45 +00:00
l = [i - 1] + module.layers # layers
2020-04-07 21:19:43 +00:00
sh = [list(x.shape)] + [list(out[i].shape) for i in module.layers] # shapes
str = ' >> ' + ' + '.join(['layer %g %s' % x for x in zip(l, sh)])
2020-04-05 21:47:41 +00:00
x = module(x, out) # WeightedFeatureFusion(), FeatureConcat()
2020-04-05 22:22:32 +00:00
elif name == 'YOLOLayer':
2020-04-14 01:25:59 +00:00
yolo_out.append(module(x, out))
2020-04-05 20:49:13 +00:00
else: # run module directly, i.e. mtype = 'convolutional', 'upsample', 'maxpool', 'batchnorm2d' etc.
x = module(x)
2020-03-19 23:41:42 +00:00
out.append(x if self.routs[i] else [])
2020-01-23 21:52:17 +00:00
if verbose:
2020-04-05 22:22:32 +00:00
print('%g/%g %s -' % (i, len(self.module_list), name), list(x.shape), str)
2020-02-20 02:26:45 +00:00
str = ''
2018-08-26 08:51:39 +00:00
2020-02-19 20:59:56 +00:00
if self.training: # train
2020-02-19 23:16:00 +00:00
return yolo_out
2020-02-19 20:59:56 +00:00
elif ONNX_EXPORT: # export
2020-02-19 23:16:00 +00:00
x = [torch.cat(x, 0) for x in zip(*yolo_out)]
2020-01-11 21:11:30 +00:00
return x[0], torch.cat(x[1:3], 1) # scores, boxes: 3780x80, 3780x4
else: # inference or test
x, p = zip(*yolo_out) # inference output, training output
x = torch.cat(x, 1) # cat yolo outputs
if augment: # de-augment results
x = torch.split(x, nb, dim=0)
2020-04-07 21:19:43 +00:00
x[1][..., :4] /= s[0] # scale
x[1][..., 0] = img_size[1] - x[1][..., 0] # flip lr
2020-04-07 21:19:43 +00:00
x[2][..., :4] /= s[1] # scale
x = torch.cat(x, 1)
return x, p
2018-08-26 08:51:39 +00:00
2019-04-20 20:46:23 +00:00
def fuse(self):
# Fuse Conv2d + BatchNorm2d layers throughout model
print('Fusing layers...')
2019-04-20 20:46:23 +00:00
fused_list = nn.ModuleList()
for a in list(self.children())[0]:
2019-08-09 10:44:47 +00:00
if isinstance(a, nn.Sequential):
for i, b in enumerate(a):
if isinstance(b, nn.modules.batchnorm.BatchNorm2d):
# fuse this bn layer with the previous conv2d layer
conv = a[i - 1]
fused = torch_utils.fuse_conv_and_bn(conv, b)
a = nn.Sequential(fused, *list(a.children())[i + 1:])
break
2019-04-20 20:46:23 +00:00
fused_list.append(a)
self.module_list = fused_list
2020-04-15 19:12:59 +00:00
self.info() if not ONNX_EXPORT else None # yolov3-spp reduced from 225 to 152 layers
2020-03-14 23:46:54 +00:00
def info(self, verbose=False):
torch_utils.model_info(self, verbose)
2019-04-20 20:46:23 +00:00
2018-08-26 08:51:39 +00:00
2020-04-14 00:56:12 +00:00
def get_yolo_layers(model):
return [i for i, m in enumerate(model.module_list) if m.__class__.__name__ == 'YOLOLayer'] # [89, 101, 113]
2019-02-08 21:43:05 +00:00
def load_darknet_weights(self, weights, cutoff=-1):
# Parses and loads the weights stored in 'weights'
2019-02-08 15:50:48 +00:00
2019-09-19 16:05:04 +00:00
# Establish cutoffs (load layers between 0 and cutoff. if cutoff = -1 all are loaded)
file = Path(weights).name
2019-07-29 21:37:12 +00:00
if file == 'darknet53.conv.74':
2018-10-30 13:58:26 +00:00
cutoff = 75
2019-07-29 21:37:12 +00:00
elif file == 'yolov3-tiny.conv.15':
2019-02-21 14:57:18 +00:00
cutoff = 15
2018-08-26 08:51:39 +00:00
2019-06-05 11:49:56 +00:00
# Read weights file
2019-04-23 14:48:47 +00:00
with open(weights, 'rb') as f:
2019-06-05 11:49:56 +00:00
# Read Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
self.version = np.fromfile(f, dtype=np.int32, count=3) # (int32) version info: major, minor, revision
self.seen = np.fromfile(f, dtype=np.int64, count=1) # (int64) number of images seen during training
2018-08-26 08:51:39 +00:00
2019-11-15 01:48:06 +00:00
weights = np.fromfile(f, dtype=np.float32) # the rest are weights
2018-08-26 08:51:39 +00:00
ptr = 0
2019-08-03 12:49:38 +00:00
for i, (mdef, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
if mdef['type'] == 'convolutional':
2020-02-19 22:57:58 +00:00
conv = module[0]
2019-08-03 12:49:38 +00:00
if mdef['batch_normalize']:
2018-08-26 08:51:39 +00:00
# Load BN bias, weights, running mean and running variance
2020-02-19 22:57:58 +00:00
bn = module[1]
nb = bn.bias.numel() # number of biases
2018-08-26 08:51:39 +00:00
# Bias
2020-02-19 22:57:58 +00:00
bn.bias.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.bias))
ptr += nb
2018-08-26 08:51:39 +00:00
# Weight
2020-02-19 22:57:58 +00:00
bn.weight.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.weight))
ptr += nb
2018-08-26 08:51:39 +00:00
# Running Mean
2020-02-19 22:57:58 +00:00
bn.running_mean.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.running_mean))
ptr += nb
2018-08-26 08:51:39 +00:00
# Running Var
2020-02-19 22:57:58 +00:00
bn.running_var.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.running_var))
ptr += nb
2018-08-26 08:51:39 +00:00
else:
# Load conv. bias
2020-02-19 22:57:58 +00:00
nb = conv.bias.numel()
conv_b = torch.from_numpy(weights[ptr:ptr + nb]).view_as(conv.bias)
conv.bias.data.copy_(conv_b)
ptr += nb
2018-08-26 08:51:39 +00:00
# Load conv. weights
2020-02-19 22:57:58 +00:00
nw = conv.weight.numel() # number of weights
conv.weight.data.copy_(torch.from_numpy(weights[ptr:ptr + nw]).view_as(conv.weight))
ptr += nw
2018-08-26 08:51:39 +00:00
2019-04-23 14:48:47 +00:00
def save_weights(self, path='model.weights', cutoff=-1):
# Converts a PyTorch model to Darket format (*.pt to *.weights)
# Note: Does not work if model.fuse() is applied
with open(path, 'wb') as f:
2019-06-05 11:49:56 +00:00
# Write Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
self.version.tofile(f) # (int32) version info: major, minor, revision
self.seen.tofile(f) # (int64) number of images seen during training
2019-04-23 14:48:47 +00:00
# Iterate through layers
2019-08-03 12:49:38 +00:00
for i, (mdef, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
if mdef['type'] == 'convolutional':
2019-04-23 14:48:47 +00:00
conv_layer = module[0]
# If batch norm, load bn first
2019-08-03 12:49:38 +00:00
if mdef['batch_normalize']:
2019-04-23 14:48:47 +00:00
bn_layer = module[1]
bn_layer.bias.data.cpu().numpy().tofile(f)
bn_layer.weight.data.cpu().numpy().tofile(f)
bn_layer.running_mean.data.cpu().numpy().tofile(f)
bn_layer.running_var.data.cpu().numpy().tofile(f)
# Load conv bias
else:
conv_layer.bias.data.cpu().numpy().tofile(f)
# Load conv weights
conv_layer.weight.data.cpu().numpy().tofile(f)
2018-08-26 08:51:39 +00:00
2019-04-23 14:48:47 +00:00
def convert(cfg='cfg/yolov3-spp.cfg', weights='weights/yolov3-spp.weights'):
# Converts between PyTorch and Darknet format per extension (i.e. *.weights convert to *.pt and vice versa)
# from models import *; convert('cfg/yolov3-spp.cfg', 'weights/yolov3-spp.weights')
# Initialize model
model = Darknet(cfg)
# Load weights and save
if weights.endswith('.pt'): # if PyTorch format
model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
save_weights(model, path='converted.weights', cutoff=-1)
print("Success: converted '%s' to 'converted.weights'" % weights)
elif weights.endswith('.weights'): # darknet format
_ = load_darknet_weights(model, weights)
2019-07-08 17:26:46 +00:00
chkpt = {'epoch': -1,
'best_fitness': None,
'training_results': None,
'model': model.state_dict(),
'optimizer': None}
2019-04-23 14:48:47 +00:00
torch.save(chkpt, 'converted.pt')
print("Success: converted '%s' to 'converted.pt'" % weights)
else:
print('Error: extension not supported.')
2019-09-19 16:05:04 +00:00
def attempt_download(weights):
# Attempt to download pretrained weights if not found locally
weights = weights.strip()
2019-12-06 21:44:13 +00:00
msg = weights + ' missing, try downloading from https://drive.google.com/open?id=1LezFG5g3BCW6iYaV89B2i64cqEUZD7e0'
2019-09-19 16:05:04 +00:00
if len(weights) > 0 and not os.path.isfile(weights):
2019-12-06 20:56:22 +00:00
d = {'yolov3-spp.weights': '16lYS4bcIdM2HdmyJBVDOvt3Trx6N3W2R',
'yolov3.weights': '1uTlyDWlnaqXcsKOktP5aH_zRDbfcDp-y',
'yolov3-tiny.weights': '1CCF-iNIIkYesIDzaPvdwlcf7H9zSsKZQ',
'yolov3-spp.pt': '1f6Ovy3BSq2wYq4UfvFUpxJFNDFfrIDcR',
'yolov3.pt': '1SHNFyoe5Ni8DajDNEqgB2oVKBb_NoEad',
'yolov3-tiny.pt': '10m_3MlpQwRtZetQxtksm9jqHrPTHZ6vo',
'darknet53.conv.74': '1WUVBid-XuoUBmvzBVUCBl_ELrzqwA8dJ',
'yolov3-tiny.conv.15': '1Bw0kCpplxUqyRYAJr9RY9SGnOJbo9nEj',
2020-02-17 07:12:07 +00:00
'yolov3-spp-ultralytics.pt': '1UcR-zVoMs7DH5dj3N1bswkiQTA4dmKF4'}
2019-12-06 20:56:22 +00:00
2019-12-06 21:47:17 +00:00
file = Path(weights).name
2019-12-06 21:44:13 +00:00
if file in d:
r = gdrive_download(id=d[file], name=weights)
2019-12-06 20:56:22 +00:00
else: # download from pjreddie.com
2019-12-06 21:44:13 +00:00
url = 'https://pjreddie.com/media/files/' + file
print('Downloading ' + url)
r = os.system('curl -f ' + url + ' -o ' + weights)
2019-12-06 21:50:16 +00:00
# Error check
if not (r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6): # weights exist and > 1MB
os.system('rm ' + weights) # remove partial downloads
raise Exception(msg)