def __init__(self, config, model_path):
"""
Initialization of PytorchYolo
:param config: vision config dict
:param model_path: path to the yolo model
"""
weightpath = os.path.join(model_path, "yolo_weights.weights")
configpath = os.path.join(model_path, "config.cfg")
self.model = torch_models.load_model(configpath, weightpath)
self._image = None
super(YoloHandlerPytorch, self).__init__(config, model_path)
def evaluate_model_file(model_path,
weights_path,
img_path,
class_names,
batch_size=8,
img_size=416,
n_cpu=8,
iou_thres=0.5,
conf_thres=0.5,
nms_thres=0.5,
verbose=True):
"""Evaluate model on validation dataset.
:param model_path: Path to model definition file (.cfg)
:type model_path: str
:param weights_path: Path to weights or checkpoint file (.weights or .pth)
:type weights_path: str
:param img_path: Path to file containing all paths to validation images.
:type img_path: str
:param class_names: List of class names
:type class_names: [str]
:param batch_size: Size of each image batch, defaults to 8
:type batch_size: int, optional
:param img_size: Size of each image dimension for yolo, defaults to 416
:type img_size: int, optional
:param n_cpu: Number of cpu threads to use during batch generation, defaults to 8
:type n_cpu: int, optional
:param iou_thres: IOU threshold required to qualify as detected, defaults to 0.5
:type iou_thres: float, optional
:param conf_thres: Object confidence threshold, defaults to 0.5
:type conf_thres: float, optional
:param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
:type nms_thres: float, optional
:param verbose: If True, prints stats of model, defaults to True
:type verbose: bool, optional
:return: Returns precision, recall, AP, f1, ap_class
"""
dataloader = _create_validation_data_loader(img_path, batch_size, img_size,
n_cpu)
model = load_model(model_path, weights_path)
metrics_output = _evaluate(model, dataloader, class_names, img_size,
iou_thres, conf_thres, nms_thres, verbose)
return metrics_output
def detect_directory(model_path,
weights_path,
img_path,
classes,
output_path,
batch_size=8,
img_size=416,
n_cpu=8,
conf_thres=0.5,
nms_thres=0.5):
"""Detects objects on all images in specified directory and saves output images with drawn detections.
:param model_path: Path to model definition file (.cfg)
:type model_path: str
:param weights_path: Path to weights or checkpoint file (.weights or .pth)
:type weights_path: str
:param img_path: Path to directory with images to inference
:type img_path: str
:param classes: List of class names
:type classes: [str]
:param output_path: Path to output directory
:type output_path: str
:param batch_size: Size of each image batch, defaults to 8
:type batch_size: int, optional
:param img_size: Size of each image dimension for yolo, defaults to 416
:type img_size: int, optional
:param n_cpu: Number of cpu threads to use during batch generation, defaults to 8
:type n_cpu: int, optional
:param conf_thres: Object confidence threshold, defaults to 0.5
:type conf_thres: float, optional
:param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
:type nms_thres: float, optional
"""
dataloader = _create_data_loader(img_path, batch_size, img_size, n_cpu)
model = load_model(model_path, weights_path)
img_detections, imgs = detect(model, dataloader, output_path, img_size,
conf_thres, nms_thres)
_draw_and_save_output_images(img_detections, imgs, img_size, output_path,
classes)
def run():
print_environment_info()
parser = argparse.ArgumentParser(description="Trains the YOLO model.")
parser.add_argument("-m",
"--model",
type=str,
default="config/yolov3.cfg",
help="Path to model definition file (.cfg)")
parser.add_argument("-d",
"--data",
type=str,
default="config/coco.data",
help="Path to data config file (.data)")
parser.add_argument("-e",
"--epochs",
type=int,
default=300,
help="Number of epochs")
parser.add_argument("-v",
"--verbose",
action='store_true',
help="Makes the training more verbose")
parser.add_argument(
"--n_cpu",
type=int,
default=8,
help="Number of cpu threads to use during batch generation")
parser.add_argument(
"--pretrained_weights",
type=str,
help=
"Path to checkpoint file (.weights or .pth). Starts training from checkpoint model"
)
parser.add_argument("--checkpoint_interval",
type=int,
default=1,
help="Interval of epochs between saving model weights")
parser.add_argument(
"--evaluation_interval",
type=int,
default=1,
help="Interval of epochs between evaluations on validation set")
parser.add_argument("--multiscale_training",
action="store_false",
help="Allow for multi-scale training")
parser.add_argument(
"--iou_thres",
type=float,
default=0.5,
help="Evaluation: IOU threshold required to qualify as detected")
parser.add_argument("--conf_thres",
type=float,
default=0.1,
help="Evaluation: Object confidence threshold")
parser.add_argument(
"--nms_thres",
type=float,
default=0.5,
help="Evaluation: IOU threshold for non-maximum suppression")
parser.add_argument(
"--logdir",
type=str,
default="logs",
help="Directory for training log files (e.g. for TensorBoard)")
parser.add_argument("--seed",
type=int,
default=-1,
help="Makes results reproducable. Set -1 to disable.")
args = parser.parse_args()
print(f"Command line arguments: {args}")
if args.seed != -1:
provide_determinism(args.seed)
logger = Logger(args.logdir) # Tensorboard logger
# Create output directories if missing
os.makedirs("output", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)
# Get data configuration
data_config = parse_data_config(args.data)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ############
# Create model
# ############
model = load_model(args.model, args.pretrained_weights)
# Print model
if args.verbose:
summary(model,
input_size=(3, model.hyperparams['height'],
model.hyperparams['height']))
mini_batch_size = model.hyperparams['batch'] // model.hyperparams[
'subdivisions']
# #################
# Create Dataloader
# #################
# Load training dataloader
dataloader = _create_data_loader(train_path, mini_batch_size,
model.hyperparams['height'], args.n_cpu,
args.multiscale_training)
# Load validation dataloader
validation_dataloader = _create_validation_data_loader(
valid_path, mini_batch_size, model.hyperparams['height'], args.n_cpu)
# ################
# Create optimizer
# ################
params = [p for p in model.parameters() if p.requires_grad]
if (model.hyperparams['optimizer'] in [None, "adam"]):
optimizer = optim.Adam(
params,
lr=model.hyperparams['learning_rate'],
weight_decay=model.hyperparams['decay'],
)
elif (model.hyperparams['optimizer'] == "sgd"):
optimizer = optim.SGD(params,
lr=model.hyperparams['learning_rate'],
weight_decay=model.hyperparams['decay'],
momentum=model.hyperparams['momentum'])
else:
print("Unknown optimizer. Please choose between (adam, sgd).")
for epoch in range(args.epochs):
print("\n---- Training Model ----")
model.train() # Set model to training mode
for batch_i, (_, imgs, targets) in enumerate(
tqdm.tqdm(dataloader, desc=f"Training Epoch {epoch}")):
batches_done = len(dataloader) * epoch + batch_i
imgs = imgs.to(device, non_blocking=True)
targets = targets.to(device)
outputs = model(imgs)
loss, loss_components = compute_loss(outputs, targets, model)
loss.backward()
###############
# Run optimizer
###############
if batches_done % model.hyperparams['subdivisions'] == 0:
# Adapt learning rate
# Get learning rate defined in cfg
lr = model.hyperparams['learning_rate']
if batches_done < model.hyperparams['burn_in']:
# Burn in
lr *= (batches_done / model.hyperparams['burn_in'])
else:
# Set and parse the learning rate to the steps defined in the cfg
for threshold, value in model.hyperparams['lr_steps']:
if batches_done > threshold:
lr *= value
# Log the learning rate
logger.scalar_summary("train/learning_rate", lr, batches_done)
# Set learning rate
for g in optimizer.param_groups:
g['lr'] = lr
# Run optimizer
optimizer.step()
# Reset gradients
optimizer.zero_grad()
# ############
# Log progress
# ############
if args.verbose:
print(
AsciiTable([
["Type", "Value"],
["IoU loss", float(loss_components[0])],
["Object loss",
float(loss_components[1])],
["Class loss", float(loss_components[2])],
["Loss", float(loss_components[3])],
["Batch loss", to_cpu(loss).item()],
]).table)
# Tensorboard logging
tensorboard_log = [("train/iou_loss", float(loss_components[0])),
("train/obj_loss", float(loss_components[1])),
("train/class_loss", float(loss_components[2])),
("train/loss", to_cpu(loss).item())]
logger.list_of_scalars_summary(tensorboard_log, batches_done)
model.seen += imgs.size(0)
# #############
# Save progress
# #############
# Save model to checkpoint file
if epoch % args.checkpoint_interval == 0:
checkpoint_path = f"checkpoints/yolov3_ckpt_{epoch}.pth"
print(f"---- Saving checkpoint to: '{checkpoint_path}' ----")
torch.save(model.state_dict(), checkpoint_path)
# ########
# Evaluate
# ########
if epoch % args.evaluation_interval == 0:
print("\n---- Evaluating Model ----")
# Evaluate the model on the validation set
metrics_output = _evaluate(model,
validation_dataloader,
class_names,
img_size=model.hyperparams['height'],
iou_thres=args.iou_thres,
conf_thres=args.conf_thres,
nms_thres=args.nms_thres,
verbose=args.verbose)
if metrics_output is not None:
precision, recall, AP, f1, ap_class = metrics_output
evaluation_metrics = [("validation/precision",
precision.mean()),
("validation/recall", recall.mean()),
("validation/mAP", AP.mean()),
("validation/f1", f1.mean())]
logger.list_of_scalars_summary(evaluation_metrics, epoch)
import torch
from pytorchyolo import models
from nni.compression.pytorch import ModelSpeedup
from nni.algorithms.compression.pytorch.pruning import L1FilterPruner, LevelPruner
from nni.compression.pytorch.utils import not_safe_to_prune
# The Yolo can be downloaded at https://github.com/eriklindernoren/PyTorch-YOLOv3.git
prefix = '/home/user/PyTorch-YOLOv3' # replace this path with yours
# Load the YOLO model
model = models.load_model(
"%s/config/yolov3.cfg" % prefix,
"%s/yolov3.weights" % prefix).cpu()
model.eval()
dummy_input = torch.rand(8, 3, 320, 320)
model(dummy_input)
# Generate the config list for pruner
# Filter the layers that may not be able to prune
not_safe = not_safe_to_prune(model, dummy_input)
cfg_list = []
for name, module in model.named_modules():
if name in not_safe:
continue
if isinstance(module, torch.nn.Conv2d):
cfg_list.append({'op_types':['Conv2d'], 'sparsity':0.6, 'op_names':[name]})
# Prune the model
pruner = L1FilterPruner(model, cfg_list)
pruner.compress()
pruner.export_model('./model', './mask')
pruner._unwrap_model()
# Speedup the model
from pytorchyolo.utils.loss import compute_loss
from pytorchyolo.models import load_model
from pytorchyolo.models import Darknet
model = Net()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
learning_rate = 1e-6
params = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.Adam(lr=model.hyperparams['learning rate'], params=params)
epochs = 10
for epoch in range(epochs):
print('training model....')
model = load_model(model_path='yolo_model.py')
model.train()
for batch_i, (_, imgs, targets) in enumerate(tqdm.tqdm(dataloader, desc=f"Training Epoch{epoch}")):
batches_done = len(dataloader) * epoch + batch_i
imgs = imgs.to(device, non_blocking=True)
targets = targets.to(device)
outputs = model(imgs)
loss, loss_components = compute_loss(outputs, targets, model)
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(epoch+1)
