rename /checkpoints to /weights
This commit is contained in:
parent
553254bbd6
commit
0ae90d0fb7
|
@ -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
|
||||||
|
|
2
test.py
2
test.py
|
@ -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')
|
||||||
|
|
12
train.py
12
train.py
|
@ -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
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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')
|
||||||
|
|
Loading…
Reference in New Issue