diff --git a/models.py b/models.py index 877e6439..3968b726 100755 --- a/models.py +++ b/models.py @@ -172,15 +172,19 @@ class YOLOLayer(nn.Module): class Darknet(nn.Module): """YOLOv3 object detection model""" - def __init__(self, cfg_path, img_size=416): + def __init__(self, cfg, img_size=(416, 416)): super(Darknet, self).__init__() - self.module_defs = parse_model_cfg(cfg_path) - self.module_defs[0]['cfg'] = cfg_path + self.module_defs = parse_model_cfg(cfg) + self.module_defs[0]['cfg'] = cfg self.module_defs[0]['height'] = img_size self.hyperparams, self.module_list = create_modules(self.module_defs) self.yolo_layers = get_yolo_layers(self) + # Needed to write header when saving weights + self.header_info = np.zeros(5, dtype=np.int32) # First five are header values + self.seen = self.header_info[3] # number of images seen during training + def forward(self, x, var=None): img_size = max(x.shape[-2:]) layer_outputs = [] @@ -270,15 +274,14 @@ def load_darknet_weights(self, weights, cutoff=-1): cutoff = 15 # Open the weights file - fp = open(weights, 'rb') - header = np.fromfile(fp, dtype=np.int32, count=5) # First five are header values + with open(weights, 'rb') as f: + header = np.fromfile(f, dtype=np.int32, count=5) # First five are header values - # Needed to write header when saving weights - self.header_info = header + # Needed to write header when saving weights + self.header_info = header - self.seen = header[3] # number of images seen during training - weights = np.fromfile(fp, dtype=np.float32) # The rest are weights - fp.close() + self.seen = header[3] # number of images seen during training + weights = np.fromfile(f, dtype=np.float32) # The rest are weights ptr = 0 for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): @@ -319,26 +322,49 @@ def load_darknet_weights(self, weights, cutoff=-1): return cutoff -def save_weights(self, path, cutoff=-1): - fp = open(path, 'wb') - self.header_info[3] = self.seen # number of images seen during training - self.header_info.tofile(fp) +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: + self.header_info[3] = self.seen # number of images seen during training + self.header_info.tofile(f) - # Iterate through layers - for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): - if module_def['type'] == 'convolutional': - conv_layer = module[0] - # If batch norm, load bn first - if module_def['batch_normalize']: - bn_layer = module[1] - bn_layer.bias.data.cpu().numpy().tofile(fp) - bn_layer.weight.data.cpu().numpy().tofile(fp) - bn_layer.running_mean.data.cpu().numpy().tofile(fp) - bn_layer.running_var.data.cpu().numpy().tofile(fp) - # Load conv bias - else: - conv_layer.bias.data.cpu().numpy().tofile(fp) - # Load conv weights - conv_layer.weight.data.cpu().numpy().tofile(fp) + # Iterate through layers + for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): + if module_def['type'] == 'convolutional': + conv_layer = module[0] + # If batch norm, load bn first + if module_def['batch_normalize']: + 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) - fp.close() + +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) + chkpt = {'epoch': -1, 'best_loss': None, 'model': model.state_dict(), 'optimizer': None} + torch.save(chkpt, 'converted.pt') + print("Success: converted '%s' to 'converted.pt'" % weights) + + else: + print('Error: extension not supported.')