def val(self, epoch):
self.model.eval()
self.ema.ema.eval()
predict_list = list()
target_list = list()
if self.local_rank == 0:
pbar = tqdm(self.vloader)
else:
pbar = self.vloader
for img_tensor, targets_tensor in pbar:
_, _, h, w = img_tensor.shape
targets_tensor[:, 1:] = targets_tensor[:, 1:] * torch.tensor(data=[w, h, w, h])
targets_tensor[:, [1, 2]] = targets_tensor[:, [1, 2]] - targets_tensor[:, [3, 4]] * 0.5
targets_tensor[:, [3, 4]] = targets_tensor[:, [1, 2]] + targets_tensor[:, [3, 4]]
img_tensor = img_tensor.to(self.device)
targets_tensor = targets_tensor.to(self.device)
predicts = self.ema.ema(img_tensor)['predicts']
for i, pred in enumerate(predicts):
if pred is not None:
pred = torch.cat([pred, torch.zeros_like(pred[..., [0]])], dim=-1)
predict_list.append(pred)
targets_sample = targets_tensor[targets_tensor[:, 0] == i][:, 1:]
targets_sample = torch.cat([torch.zeros_like(targets_sample[..., [0]]), targets_sample], dim=-1)
target_list.append(targets_sample)
mp, mr, map50, map = coco_map(predict_list, target_list)
mp = reduce_sum(torch.tensor(mp, device=self.device)).item() / self.gpu_num
mr = reduce_sum(torch.tensor(mr, device=self.device)).item() / self.gpu_num
map50 = reduce_sum(torch.tensor(map50, device=self.device)).item() / self.gpu_num
map = reduce_sum(torch.tensor(map, device=self.device)).item() / self.gpu_num
if self.local_rank == 0:
print("epoch: {:2d}|gpu_num:{:d}|mp:{:6.4f}|mr:{:6.4f}|map50:{:6.4f}|map:{:6.4f}"
.format(epoch + 1,
self.gpu_num,
mp * 100,
mr * 100,
map50 * 100,
map * 100))
last_weight_path = os.path.join(self.val_cfg['weight_path'],
"{:s}_{:s}_last.pth"
.format(self.cfg['model_name'], self.model_cfg['scale_name']))
best_map_weight_path = os.path.join(self.val_cfg['weight_path'],
"{:s}_{:s}_best_map.pth"
.format(self.cfg['model_name'], self.model_cfg['scale_name']))
ema_static = self.ema.ema.state_dict()
cpkt = {
"ema": ema_static,
"map": map * 100,
"epoch": epoch,
}
if self.local_rank != 0:
return
torch.save(cpkt, last_weight_path)
if map > self.best_map:
torch.save(cpkt, best_map_weight_path)
self.best_map = map
def val(self, epoch):
predict_list = list()
target_list = list()
self.model.eval()
self.ema.ema.eval()
if self.local_rank == 0:
pbar = tqdm(self.vloader)
else:
pbar = self.vloader
for img_tensor, valid_size, targets_tensor, mask_tensor, batch_len in pbar:
img_tensor = img_tensor.to(self.device)
targets_tensor = targets_tensor.to(self.device)
predicts = self.ema.ema(img_tensor,
valid_size=valid_size)['predicts']
predicts = [p[0] for p in predicts]
for pred, target in zip(predicts, targets_tensor.split(batch_len)):
predict_list.append(pred)
target_list.append(target)
mp, mr, map50, mean_ap = coco_map(predict_list, target_list)
mp = reduce_sum(torch.tensor(mp, device=self.device)) / self.gpu_num
mr = reduce_sum(torch.tensor(mr, device=self.device)) / self.gpu_num
map50 = reduce_sum(torch.tensor(map50,
device=self.device)) / self.gpu_num
mean_ap = reduce_sum(torch.tensor(mean_ap,
device=self.device)) / self.gpu_num
if self.local_rank == 0:
print("*" * 20, "eval start", "*" * 20)
print(
"epoch: {:2d}|mp:{:6.4f}|mr:{:6.4f}|map50:{:6.4f}|map:{:6.4f}".
format(epoch + 1, mp * 100, mr * 100, map50 * 100,
mean_ap * 100))
print("*" * 20, "eval end", "*" * 20)
last_weight_path = os.path.join(
self.val_cfg['weight_path'],
"{:s}_{:s}_last.pth".format(self.cfg['model_name'],
self.model_cfg['backbone']))
best_map_weight_path = os.path.join(
self.val_cfg['weight_path'],
"{:s}_{:s}_best_map.pth".format(self.cfg['model_name'],
self.model_cfg['backbone']))
ema_static = self.ema.ema.state_dict()
cpkt = {
"ema": ema_static,
"map": mean_ap * 100,
"epoch": epoch,
}
if self.local_rank != 0:
return
torch.save(cpkt, last_weight_path)
if mean_ap > self.best_map:
torch.save(cpkt, best_map_weight_path)
self.best_map = mean_ap
def __call__(self, cls_predicts, reg_predicts, targets, shape):
h, w = shape
shape_norm = torch.tensor([w, h, w, h], device=cls_predicts.device)
pos_num = len(targets['target'])
gt_boxes = targets['target'].split(targets['batch_len'])
if cls_predicts.dtype == torch.float16:
cls_predicts = cls_predicts.float()
all_imme_idx = list()
all_batch_idx = list()
all_proposal_idx = list()
all_cls_label_idx = list()
all_box_targets = list()
for imme_idx, batch_cls_predict, batch_reg_predict in zip(
range(len(cls_predicts)), cls_predicts, reg_predicts):
matches = self.matcher(batch_cls_predict.detach(),
batch_reg_predict.detach(), gt_boxes,
shape_norm)
match_cls_bidx = sum([[i] * len(j) for i, j, _ in matches], [])
match_proposal_idx = sum([j for _, j, _ in matches], [])
match_cls_label_idx = torch.cat(
[gt_boxes[i][:, 0][k].long() for i, _, k in matches])
match_box = torch.cat(
[gt_boxes[i][:, 1:][k] for i, _, k in matches])
all_imme_idx.append([imme_idx] * len(match_cls_bidx))
all_batch_idx.append(match_cls_bidx)
all_proposal_idx.append(match_proposal_idx)
all_cls_label_idx.append(match_cls_label_idx)
all_box_targets.append(match_box)
all_imme_idx = sum(all_imme_idx, [])
all_batch_idx = sum(all_batch_idx, [])
all_proposal_idx = sum(all_proposal_idx, [])
all_cls_label_idx = torch.cat(all_cls_label_idx)
all_box_targets = torch.cat(all_box_targets)
cls_targets = torch.zeros_like(cls_predicts)
cls_targets[all_imme_idx, all_batch_idx, all_proposal_idx,
all_cls_label_idx] = 1.0
box_pred = reg_predicts[all_imme_idx, all_batch_idx, all_proposal_idx]
cls_loss = self.cls_weights * focal_loss(cls_predicts.sigmoid(),
cls_targets).sum()
box_loss = self.iou_weights * self.iou_loss(box_pred,
all_box_targets).sum()
l1_loss = self.l1_weights * torch.nn.functional.l1_loss(
box_pred / shape_norm[None, :],
all_box_targets / shape_norm[None, :],
reduction="none").sum()
pos_num = reduce_sum(torch.tensor(
pos_num, device=cls_predicts.device)).item() / get_gpu_num_solo()
return cls_loss / pos_num, box_loss / pos_num, l1_loss / pos_num, pos_num
def train(self, epoch):
self.loss_logger.reset()
self.cls_loss_logger.reset()
self.l1_loss_logger.reset()
self.iou_loss_logger.reset()
self.match_num_logger.reset()
self.model.train()
if self.local_rank == 0:
pbar = tqdm(self.tloader)
else:
pbar = self.tloader
for i, (img_tensor, targets_tensor, batch_len) in enumerate(pbar):
_, _, h, w = img_tensor.shape
with torch.no_grad():
img_tensor = img_tensor.to(self.device)
targets_tensor = targets_tensor.to(self.device)
self.optimizer.zero_grad()
if self.scaler is not None:
with amp.autocast(enabled=True):
out = self.model(img_tensor,
targets={
"target": targets_tensor,
"batch_len": batch_len
})
cls_loss = out['cls_loss']
l1_loss = out['l1_loss']
iou_loss = out['iou_loss']
match_num = out['match_num']
loss = cls_loss + l1_loss + iou_loss
self.scaler.scale(loss).backward()
self.lr_adjuster(self.optimizer, i, epoch)
self.scaler.step(self.optimizer)
self.scaler.update()
else:
out = self.model(img_tensor,
targets={
"target": targets_tensor,
"batch_len": batch_len
})
cls_loss = out['cls_loss']
l1_loss = out['l1_loss']
iou_loss = out['iou_loss']
match_num = out['match_num']
loss = cls_loss + l1_loss + iou_loss
loss.backward()
self.lr_adjuster(self.optimizer, i, epoch)
self.optimizer.step()
self.ema.update(self.model)
lr = self.optimizer.param_groups[0]['lr']
self.loss_logger.update(loss.item())
self.iou_loss_logger.update(iou_loss.item())
self.l1_loss_logger.update(l1_loss.item())
self.cls_loss_logger.update(cls_loss.item())
self.match_num_logger.update(match_num)
str_template = \
"epoch:{:2d}|match_num:{:0>4d}|size:{:3d}|loss:{:6.4f}|cls:{:6.4f}|l1:{:6.4f}|iou:{:6.4f}|lr:{:8.6f}"
if self.local_rank == 0:
pbar.set_description(
str_template.format(epoch + 1, int(match_num), h,
self.loss_logger.avg(),
self.cls_loss_logger.avg(),
self.l1_loss_logger.avg(),
self.iou_loss_logger.avg(), lr))
self.ema.update_attr(self.model)
loss_avg = reduce_sum(
torch.tensor(self.loss_logger.avg(),
device=self.device)) / self.gpu_num
iou_loss_avg = reduce_sum(
torch.tensor(self.iou_loss_logger.avg(),
device=self.device)).item() / self.gpu_num
l1_loss_avg = reduce_sum(
torch.tensor(self.l1_loss_logger.avg(),
device=self.device)).item() / self.gpu_num
cls_loss_avg = reduce_sum(
torch.tensor(self.cls_loss_logger.avg(),
device=self.device)).item() / self.gpu_num
match_num_sum = reduce_sum(
torch.tensor(self.match_num_logger.sum(),
device=self.device)).item() / self.gpu_num
if self.local_rank == 0:
final_template = "epoch:{:2d}|match_num:{:d}|loss:{:6.4f}|cls:{:6.4f}|l1:{:6.4f}|iou:{:6.4f}"
print(
final_template.format(epoch + 1, int(match_num_sum), loss_avg,
cls_loss_avg, l1_loss_avg, iou_loss_avg))
def train(self, epoch):
self.loss.reset()
self.cls_loss.reset()
self.box_loss.reset()
self.mask_loss.reset()
self.model.train()
if self.local_rank == 0:
pbar = tqdm(self.tloader)
else:
pbar = self.tloader
for i, (img_tensor, valid_size, targets_tensor, mask_tensor, batch_len) in enumerate(pbar):
_, _, h, w = img_tensor.shape
with torch.no_grad():
img_tensor = img_tensor.to(self.device)
targets_tensor = targets_tensor.to(self.device)
mask_tensor = mask_tensor.to(self.device)
self.optimizer.zero_grad()
if self.scaler is not None:
with amp.autocast(enabled=True):
out = self.model(img_tensor, valid_size=valid_size,
targets={"target": targets_tensor,
"batch_len": batch_len,
"mask": mask_tensor})
box_loss = sum([v for k, v in out.items() if "box" in k])
cls_loss = sum([v for k, v in out.items() if "cls" in k])
mask_loss = sum([v for k, v in out.items() if "mask" in k])
loss = box_loss + cls_loss + mask_loss
self.scaler.scale(loss).backward()
self.lr_adjuster(self.optimizer, i, epoch)
self.scaler.step(self.optimizer)
self.scaler.update()
else:
out = self.model(img_tensor, valid_size=valid_size,
targets={"target": targets_tensor,
"batch_len": batch_len,
"mask": mask_tensor})
box_loss = sum([v for k, v in out.items() if "box" in k])
cls_loss = sum([v for k, v in out.items() if "cls" in k])
mask_loss = sum([v for k, v in out.items() if "mask" in k])
loss = box_loss + cls_loss + mask_loss
loss.backward()
self.lr_adjuster(self.optimizer, i, epoch)
self.optimizer.step()
self.ema.update(self.model)
lr = self.optimizer.param_groups[0]['lr']
self.loss.update(loss.item())
self.cls_loss.update(cls_loss.item())
self.box_loss.update(box_loss.item())
self.mask_loss.update(mask_loss.item())
str_template = "epoch:{:2d}|size:{:3d}|loss:{:6.4f}|cls:{:6.4f}|box:{:6.4f}|mask:{:6.4f}|lr:{:8.6f}"
if self.local_rank == 0:
pbar.set_description(str_template.format(
epoch,
h,
self.loss.avg(),
self.cls_loss.avg(),
self.box_loss.avg(),
self.mask_loss.avg(),
lr)
)
self.ema.update_attr(self.model)
loss_avg = reduce_sum(torch.tensor(self.loss.avg(), device=self.device)) / self.gpu_num
cls_avg = reduce_sum(torch.tensor(self.cls_loss.avg(), device=self.device)) / self.gpu_num
box_avg = reduce_sum(torch.tensor(self.box_loss.avg(), device=self.device)) / self.gpu_num
mask_avg = reduce_sum(torch.tensor(self.mask_loss.avg(), device=self.device)) / self.gpu_num
if self.local_rank == 0:
final_template = "epoch:{:2d}|loss:{:6.4f}|cls:{:6.4f}|box:{:6.4f}|mask:{:6.4f}"
print(final_template.format(
epoch,
loss_avg,
cls_avg,
box_avg,
mask_avg
))
def train(self, epoch):
self.loss_logger.reset()
self.cls_loss_logger.reset()
self.box_loss_logger.reset()
self.obj_loss_logger.reset()
self.match_num_logger.reset()
self.model.train()
if self.local_rank == 0:
pbar = tqdm(self.tloader)
else:
pbar = self.tloader
for i, (img_tensor, targets_tensor, batch_len) in enumerate(pbar):
_, _, h, w = img_tensor.shape
with torch.no_grad():
img_tensor = img_tensor.to(self.device)
targets_tensor = targets_tensor.to(self.device)
self.optimizer.zero_grad()
if self.scaler is not None:
with amp.autocast(enabled=True):
out = self.model(img_tensor, targets_tensor)
cls_loss = out['cls_loss']
box_loss = out['box_loss']
obj_loss = out['obj_loss']
match_num = out['match_num']
loss = cls_loss + box_loss + obj_loss
self.scaler.scale(loss).backward()
ulr, dlr, momentum = self.lr_adjuster(
self.optimizer, i, epoch)
self.scaler.step(self.optimizer)
self.scaler.update()
else:
out = self.model(img_tensor, targets_tensor)
cls_loss = out['cls_loss']
box_loss = out['box_loss']
obj_loss = out['obj_loss']
match_num = out['match_num']
loss = cls_loss + box_loss + obj_loss
loss.backward()
ulr, dlr, momentum = self.lr_adjuster(self.optimizer, i, epoch)
self.optimizer.step()
self.ema.update(self.model)
self.loss_logger.update(loss.item())
self.obj_loss_logger.update(obj_loss.item())
self.box_loss_logger.update(box_loss.item())
self.cls_loss_logger.update(cls_loss.item())
self.match_num_logger.update(match_num)
str_template = \
"epoch:{:2d}|match_num:{:0>4d}|size:{:3d}|loss:{:6.4f}|cls:{:6.4f}|box:{:6.4f}|obj:{:6.4f}|ulr:{:8.6f}|dlr:{:8.6f}"
if self.local_rank == 0:
pbar.set_description(
str_template.format(epoch, match_num, h,
self.loss_logger.avg(),
self.cls_loss_logger.avg(),
self.box_loss_logger.avg(),
self.obj_loss_logger.avg(), ulr, dlr))
self.ema.update_attr(self.model)
loss_avg = reduce_sum(
torch.tensor(self.loss_logger.avg(),
device=self.device)) / self.gpu_num
obj_loss_avg = reduce_sum(
torch.tensor(self.obj_loss_logger.avg(),
device=self.device)).item() / self.gpu_num
box_loss_avg = reduce_sum(
torch.tensor(self.box_loss_logger.avg(),
device=self.device)).item() / self.gpu_num
cls_loss_avg = reduce_sum(
torch.tensor(self.cls_loss_logger.avg(),
device=self.device)).item() / self.gpu_num
match_num_sum = reduce_sum(
torch.tensor(self.match_num_logger.sum(),
device=self.device)).item() / self.gpu_num
if self.local_rank == 0:
final_template = "epoch:{:2d}|match_num:{:d}|loss:{:6.4f}|cls:{:6.4f}|box:{:6.4f}|obj:{:6.4f}"
print(
final_template.format(epoch, int(match_num_sum), loss_avg,
cls_loss_avg, box_loss_avg,
obj_loss_avg))