rename /checkpoints to /weights

This commit is contained in:
Glenn Jocher 2018-10-27 00:42:34 +02:00
parent 553254bbd6
commit 0ae90d0fb7
6 changed files with 10 additions and 10 deletions

View File

@ -33,7 +33,7 @@ def detect(opt):
# Load model # Load model
model = Darknet(opt.cfg, opt.img_size) model = Darknet(opt.cfg, opt.img_size)
weights_path = 'checkpoints/yolov3.pt' weights_path = 'weights/yolov3.pt'
if weights_path.endswith('.weights'): # saved in darknet format if weights_path.endswith('.weights'): # saved in darknet format
load_weights(model, weights_path) load_weights(model, weights_path)
else: # endswith('.pt'), saved in pytorch format else: # endswith('.pt'), saved in pytorch format

View File

@ -8,7 +8,7 @@ parser = argparse.ArgumentParser()
parser.add_argument('-batch_size', type=int, default=32, help='size of each image batch') parser.add_argument('-batch_size', type=int, default=32, help='size of each image batch')
parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file') parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file')
parser.add_argument('-data_config_path', type=str, default='cfg/coco.data', help='path to data config file') parser.add_argument('-data_config_path', type=str, default='cfg/coco.data', help='path to data config file')
parser.add_argument('-weights_path', type=str, default='checkpoints/yolov3.pt', help='path to weights file') parser.add_argument('-weights_path', type=str, default='weights/yolov3.pt', help='path to weights file')
parser.add_argument('-class_path', type=str, default='data/coco.names', help='path to class label file') parser.add_argument('-class_path', type=str, default='data/coco.names', help='path to class label file')
parser.add_argument('-iou_thres', type=float, default=0.5, help='iou threshold required to qualify as detected') parser.add_argument('-iou_thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
parser.add_argument('-conf_thres', type=float, default=0.5, help='object confidence threshold') parser.add_argument('-conf_thres', type=float, default=0.5, help='object confidence threshold')

View File

@ -28,7 +28,7 @@ if cuda:
def main(opt): def main(opt):
os.makedirs('checkpoints', exist_ok=True) os.makedirs('weights', exist_ok=True)
# Configure run # Configure run
data_config = parse_data_config(opt.data_config_path) data_config = parse_data_config(opt.data_config_path)
@ -48,7 +48,7 @@ def main(opt):
start_epoch = 0 start_epoch = 0
best_loss = float('inf') best_loss = float('inf')
if opt.resume: if opt.resume:
checkpoint = torch.load('checkpoints/latest.pt', map_location='cpu') checkpoint = torch.load('weights/latest.pt', map_location='cpu')
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:
@ -175,15 +175,15 @@ def main(opt):
'best_loss': best_loss, 'best_loss': best_loss,
'model': model.state_dict(), 'model': model.state_dict(),
'optimizer': optimizer.state_dict()} 'optimizer': optimizer.state_dict()}
torch.save(checkpoint, 'checkpoints/latest.pt') torch.save(checkpoint, 'weights/latest.pt')
# Save best checkpoint # Save best checkpoint
if best_loss == loss_per_target: if best_loss == loss_per_target:
os.system('cp checkpoints/latest.pt checkpoints/best.pt') os.system('cp weights/latest.pt weights/best.pt')
# Save backup checkpoints every 5 epochs # Save backup weights every 5 epochs
if (epoch > 0) & (epoch % 5 == 0): if (epoch > 0) & (epoch % 5 == 0):
os.system('cp checkpoints/latest.pt checkpoints/backup' + str(epoch) + '.pt') os.system('cp weights/latest.pt weights/backup' + str(epoch) + '.pt')
# Save final model # Save final model
dt = time.time() - t0 dt = time.time() - t0

View File

@ -1,7 +1,7 @@
#!/usr/bin/env bash #!/usr/bin/env bash
# Start # Start
sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 && python3 train.py -img_size 416 -epochs 160 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 && python3 train.py -img_size 416
# Resume # Resume
python3 train.py -img_size 416 -resume 1 python3 train.py -img_size 416 -resume 1

View File

@ -410,7 +410,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
return output return output
def strip_optimizer_from_checkpoint(filename='checkpoints/best.pt'): def strip_optimizer_from_checkpoint(filename='weights/best.pt'):
# Strip optimizer from *.pt files for lighter files (reduced by 2/3 size) # Strip optimizer from *.pt files for lighter files (reduced by 2/3 size)
import torch import torch
a = torch.load(filename, map_location='cpu') a = torch.load(filename, map_location='cpu')