def __call__(self, gt_boxes, gt_labels):
boxes, labels = box_utils.assign_priors(gt_boxes, gt_labels,
self.corner_form_priors,
self.iou_threshold)
boxes = box_utils.corner_form_to_center_form(boxes)
locations = box_utils.convert_boxes_to_locations(
boxes, self.center_form_priors, self.center_variance,
self.size_variance)
return locations, labels
def __call__(self, gt_boxes, gt_labels):
if type(gt_boxes) is np.ndarray:
gt_boxes = torch.from_numpy(gt_boxes)
if type(gt_labels) is np.ndarray:
gt_labels = torch.from_numpy(gt_labels)
boxes, labels = box_utils.assign_priors(gt_boxes, gt_labels,
self.corner_form_priors, self.iou_threshold)
boxes = box_utils.corner_form_to_center_form(boxes)
locations = box_utils.convert_boxes_to_locations(boxes, self.center_form_priors, self.center_variance, self.size_variance)
return locations, labels
def train(loader,
net,
criterion,
optimizer,
device,
debug_steps=100,
epoch=-1,
is_dali=False):
net.train(True)
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
for i, data in enumerate(loader):
if not is_dali:
images, boxes, labels = data
# images = images.to(device)
# boxes = boxes.to(device)
# labels = labels.to(device)
images = images.cuda(async=True)
boxes = boxes.cuda(async=True)
labels = labels.cuda(async=True)
else:
images = data[0]["images"]
boxes = data[0]["boxes"]
labels = data[0]["labels"]
labels = labels.type(torch.cuda.LongTensor)
boxes = convert_boxes_to_locations(
boxes,
dboxes320_mv2_coco()(order="xywh"), config.center_variance,
config.size_variance)
boxes = boxes.contiguous().cuda()
optimizer.zero_grad()
confidence, locations = net(images)
regression_loss, classification_loss = criterion(
confidence, locations, labels, boxes) # TODO CHANGE BOXES
loss = regression_loss + classification_loss
loss.backward()
optimizer.step()
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
if i and i % debug_steps == 0:
lr = optimizer.param_groups[0]['lr']
avg_loss = running_loss / debug_steps
avg_reg_loss = running_regression_loss / debug_steps
avg_clf_loss = running_classification_loss / debug_steps
logging.info(f"Epoch: {epoch}, Step: {i}, LR: {lr}, " +
f"Average Loss: {avg_loss:.4f}, " +
f"Average Regression Loss {avg_reg_loss:.4f}, " +
f"Average Classification Loss: {avg_clf_loss:.4f}")
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0