def train_data(model_name, model, batch, loss_fn, device):
if model_name == "faster":
images, targets = batch
images, targets = transform_inputs(images, targets, device)
losses = model(images, targets)
loss = sum([loss for loss in losses.values()])
elif model_name == "ssd512":
images, boxes, labels = batch
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
confidence, locations = model(images)
regression_loss, classification_loss = loss_fn(confidence, locations,
labels, boxes)
loss = regression_loss + classification_loss
elif model_name == "yolov3" or model_name == "yolov3_spp" or model_name == "yolov4":
images, targets, paths, _ = batch
images = images.to(device).float() / 255.0
targets = targets.to(device)
predictions = model(images)
loss, loss_items = loss_fn(predictions, targets, model)
loss *= len(batch) / 64
return loss
def evaluate(model,data_loader,device):
"""Evaluation of the validation set
Keyword arguments:
- model: model after training with the respective weights
- data_loader: validaton set in the loader format
- device: device on which the network will be evaluated
"""
cpu_device = torch.device("cpu")
model.eval()
coco = convert_to_coco_api(data_loader.dataset)
coco_evaluator = CocoEvaluator(coco)
evaluator_times = []
proc_times = []
for image, targets in data_loader:
image, targets = transform_inputs(image, targets, device)
init = time.time()
outputs = model(image)
proc_times.append(time.time() - init)
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_times.append(time.time() - evaluator_time)
print("Averaged stats:", np.mean(evaluator_times))
print("Averaged proc time:", np.mean(proc_times))
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
return coco_evaluator
def evaluate(model, data_loader, device):
"""Evaluation of the validation set
Keyword arguments:
- model: model after training with the respective weights
- data_loader: validaton set in the loader format
- device: device on which the network will be evaluated
"""
cpu_device = torch.device("cpu")
model.eval().to(device)
coco = convert_to_coco_api(data_loader.dataset)
coco_evaluator = CocoEvaluator(coco)
evaluator_times = []
proc_times = []
for images, targets in data_loader:
res = {}
images, targets = transform_inputs(images, targets, device)
images_eval = [image.float() / 255 for image in images] #normalized
batch = torch.stack(images_eval)
init = time.time()
inf_out, eval_out = model(batch)
proc_times.append(time.time() - init)
output = non_max_suppression(inf_out, conf_thres=0.001, iou_thres=0.6)
for si, pred in enumerate(output):
height, width = images[si].shape[1:]
if pred is None:
box = torch.tensor([[0, 0, 0, 0]])
res.update({
targets[si]["image_id"].item(): {
"boxes": box,
"labels": torch.tensor([1]),
"scores": torch.tensor([0])
}
})
else:
clip_coords(pred, (height, width))
box = pred[:, :4].clone()
res.update({
targets[si]["image_id"].item(): {
"boxes": box,
"labels": pred[:, 5],
"scores": pred[:, 4]
}
})
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_times.append(time.time() - evaluator_time)
print("Averaged stats:", np.mean(evaluator_times))
print("Averaged proc time:", np.mean(proc_times))
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
return coco_evaluator
def test_data(model_name, model, batch, device):
if model_name == "faster":
images, targets = batch
images, targets = transform_inputs(images, targets, device)
images_model = copy.deepcopy(images)
torch.cuda.synchronize()
with torch.no_grad():
outputs = model(images_model)
outputs = [{k: v.to(device) for k, v in t.items()} for t in outputs]
res = {
target["image_id"].item(): output
for target, output in zip(targets, outputs)
}
elif model_name == "ssd512":
from object_detection.utils.ssd import ssd512_config as config
images, targets = batch
images_model = copy.deepcopy(images)
candidate_size = 50
sigma = 0.5
predictor = Predictor(model,
config.image_size,
config.image_mean,
config.image_std,
iou_threshold=config.iou_threshold,
candidate_size=candidate_size,
sigma=sigma,
device=device)
boxes, labels, probs = predictor.predict(images_model, 10, 0.2)
if boxes.size()[0] == 0:
outputs = {
"boxes": torch.tensor([[0, 0, 0, 0]]),
"labels": torch.tensor([0]),
"scores": torch.tensor([0])
}
else:
outputs = {"boxes": boxes, "labels": labels, "scores": probs}
res = {targets['image_id'].item(): outputs}
elif model_name == "yolov3" or model_name == "yolov3_spp" or model_name == "yolov4":
images, targets = batch
images, targets = transform_inputs(images, targets, device)
images_model = copy.deepcopy(images)
images_model = [image.float() / 255 for image in images_model]
batch_imgs = torch.stack(images_model)
labels_ = []
scores = []
res = {}
nb, _, width, height = batch_imgs.shape
torch.cuda.synchronize()
with torch.no_grad():
inf_out, train_out = model(batch_imgs)
output = non_max_suppression(
inf_out, conf_thres=0.25,
iou_thres=0.6) #conf = 0.25 to decrease the training time
for si, pred in enumerate(output):
if pred is None:
res.update({
targets[si]["image_id"].item(): {
"boxes": torch.tensor([[0, 0, 0, 0]]),
"labels": torch.tensor([1]),
"scores": torch.tensor([0])
}
})
else:
clip_coords(pred, (height, width))
box = pred[:, :4].clone()
res.update({
targets[si]["image_id"].item(): {
"boxes": box,
"labels": pred[:, 5],
"scores": pred[:, 4]
}
})
images_model = outputs = None
return images, targets, res