def ctdet_decode(hmap, regs, w_h_, K=100):
batch, cat, height, width = hmap.shape
# hmap = torch.sigmoid(hmap)
# if flip test
# if batch > 1:
# hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
# w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
# regs = regs[0:1]
#
# batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([
xs - w_h_[..., 0:1] / 2, ys - w_h_[..., 1:2] / 2,
xs + w_h_[..., 0:1] / 2, ys + w_h_[..., 1:2] / 2
],
dim=2)
detections = torch.cat([bboxes, scores, clses], dim=2)
return detections
def train_epoch(self, data_counts, data_loader, eval_loder, epoch, n_epochs):
with tqdm.tqdm(total=data_counts, desc=f'Epoch {epoch}/{n_epochs}', unit='img', ncols=150) as pbar:
step = 0
for batch in data_loader:
step += 1
load_t0 = time.time()
for k in batch:
batch[k] = batch[k].to(device=self.cfg.device, non_blocking=True)
outputs = self.model(batch['image'])
hmap, regs, w_h_ = zip(*outputs)
regs = [_tranpose_and_gather_feature(r, batch['inds']) for r in regs]
w_h_ = [_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
loss = hmap_loss + 1 * reg_loss + 0.1 * w_h_loss
self.model.zero_grad()
loss.backward()
self.optim.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
pbar.set_postfix(**{'hmap_loss': hmap_loss.item(),
'reg_loss': reg_loss.item(),
'w_h_loss': w_h_loss.item(),
'LR': self.optim.param_groups[0]['lr'],
'Batchtime': batch_time})
pbar.update(batch['image'].shape[0])
cons_acc = self._evaluate(eval_loder)
return cons_acc
def train(epoch):
print('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
outputs = model(batch['image'])
hmap, regs, w_h_, real_, imag_ = zip(*outputs)
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
real_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in real_
]
imag_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in imag_
]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
real_loss = norm_reg_loss(real_, batch['real'], batch['ind_masks'])
imag_loss = norm_reg_loss(imag_, batch['imaginary'],
batch['ind_masks'])
loss = hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.code_loss_weight * (
real_loss + imag_loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss= %.3f reg_loss= %.3f offset_loss= %.3f real_loss= %.3f imag_loss= %.3f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
real_loss.item(), imag_loss.item()) +
' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('offset_loss', w_h_loss.item(), step)
summary_writer.add_scalar('real_loss', real_loss.item(), step)
summary_writer.add_scalar('imag_loss', imag_loss.item(), step)
return
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
outputs = model(batch['image'])
hmap, regs, w_h_, codes_, shapes_ = zip(*outputs)
# print('Before gather feature: ', len(regs))
# for c in regs:
# print('regs sizes: ', c.size())
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
codes_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes_
]
shapes_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in shapes_
]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
cmm_loss = _reg_loss(shapes_, batch['shapes'], batch['ind_masks'])
# cmm_loss = contour_mapping_loss(codes_, shapes_, batch['shapes'], batch['ind_masks'])
# cmm_loss = norm_contour_mapping_loss(codes_, shapes_, batch['shapes'], batch['w_h_'], batch['ind_masks'])
loss = hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.cmm_loss_weight * cmm_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss= %.3f reg_loss= %.3f w_h_loss= %.3f cmm_loss= %.3f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
cmm_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('cmm_loss', cmm_loss.item(), step)
return
def train(epoch):
print('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
outputs = model(batch['image'])
hmap, regs, w_h_, pxpy = zip(*outputs)
# batch * C（channel） * W * H
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
pxpy = [
_tranpose_and_gather_feature(r, batch['inds']) for r in pxpy
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
# batch * K * C= batch * 128 *2
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _SmoothL1Loss(regs, batch['regs'], batch['ind_masks'])
pxpy_loss = _reg_loss(pxpy, batch['pxpy'], batch['ind_masks'])
w_h_loss = _SmoothL1Loss(w_h_, batch['w_h_'], batch['ind_masks'])
loss = hmap_loss + 10 * reg_loss + 0.1 * w_h_loss + 0.1 * pxpy_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss= %.5f reg_loss= %.5f w_h_loss= %.5f pxpy_loss= %.5f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
pxpy_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('pxpy_loss', pxpy_loss.item(), step)
return
def train(epoch):
print('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
#batch[k] = batch[k].to(device=cfg.device, non_blocking=True)
batch[k] = batch[k]
outputs = model(batch['image'].to(cfg.device))
hmap, regs, w_h_, theta = zip(*outputs)
regs = [
_tranpose_and_gather_feature(r.cpu(), batch['inds'])
for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r.cpu(), batch['inds'])
for r in w_h_
]
theta = [
_tranpose_and_gather_feature(r.cpu(), batch['inds'])
for r in theta
]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
theta_loss = _theta_loss(theta, batch['theta'], batch['ind_masks'])
loss = hmap_loss + 1. * reg_loss + 0.1 * w_h_loss + 1. * theta_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss= %.5f reg_loss= %.5f w_h_loss= %.5f theta_loss= %.5f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
theta_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
return
def ctdet_decode(hmap, regs, w_h_, K=100):
'''
hmap提取中心点位置为xs,ys
regs保存的是偏置，需要加在xs,ys上，代表精确的中心位置
w_h_保存的是对应目标的宽和高
'''
# dets = ctdet_decode(*output, K=cfg.test_topk)
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap) # 归一化到0-1
# if flip test
if batch > 1: # batch > 1代表使用了翻转
# img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
batch = 1
# 这里的nms和带anchor的目标检测方法中的不一样，这里使用的是3x3的maxpool筛选
hmap = _nms(hmap) # perform nms on heatmaps
# 找到前K个极大值点代表存在目标
scores, inds, clses, ys, xs = _topk(hmap, K=K)
# from [bs c h w] to [bs, h, w, c]
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
# xs,ys是中心坐标，w_h_[...,0:1]是w,1:2是h
bboxes = torch.cat([
xs - w_h_[..., 0:1] / 2, ys - w_h_[..., 1:2] / 2,
xs + w_h_[..., 0:1] / 2, ys + w_h_[..., 1:2] / 2
],
dim=2)
detections = torch.cat([bboxes, scores, clses], dim=2)
return detections
def train(epoch):
print('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device, non_blocking=True)
dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device, non_blocking=True)
dict_tensor.requires_grad=False
outputs = model(batch['image'])
hmap, regs, w_h_, codes_, offsets_ = zip(*outputs)
regs = [_tranpose_and_gather_feature(r, batch['inds']) for r in regs]
w_h_ = [_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_]
codes_ = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_]
offsets_ = [_tranpose_and_gather_feature(r, batch['inds']) for r in offsets_]
shapes_ = [torch.matmul(c, dict_tensor) + o for c in codes_ for o in offsets_]
# shapes_ = torch.matmul(codes_, dict_tensor) + offsets_
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
codes_loss = norm_reg_loss(codes_, batch['codes'], batch['ind_masks'])
shapes_loss = contour_mapping_loss(codes_, shapes_, batch['shapes'], batch['ind_masks'], roll=False)
loss = 2 * hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.code_loss_weight * codes_loss + \
cfg.shape_loss_weight * shapes_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print('[%d/%d-%d/%d] ' % (epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss= %.3f reg_loss= %.3f offset_loss= %.3f code_loss= %.3f shape_loss= %.3f' %
(hmap_loss.item(), reg_loss.item(), w_h_loss.item(), codes_loss.item(), shapes_loss.item()) +
' (%d samples/sec)' % (cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('offset_loss', w_h_loss.item(), step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
summary_writer.add_scalar('shape_loss', shapes_loss.item(), step)
return
def ctdet_decode(hmap, regs, w_h_, pxpy, K=100):
batch, cat, height, width = hmap.shape # C,W和H
# height , width = 128
hmap = torch.sigmoid(hmap)
# 这里，test的batch是 1
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2 # w_h_ 第一列是宽度，第二列是高度。
regs = regs[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
pxpy = _tranpose_and_gather_feature(pxpy, inds)
pxpy = pxpy.view(batch, K, 2)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
width1 = torch.abs(torch.mul(pxpy[..., 0:1], torch.cos(pxpy[...,
1:2]))) # 半宽度
height1 = torch.abs(torch.mul(pxpy[..., 0:1], torch.sin(pxpy[...,
1:2]))) # 半高度
width1 = 0.1 * width1 + 0.9 * w_h_[..., 0:1] / 2
height1 = 0.1 * height1 + 0.9 * w_h_[..., 1:2] / 2
bboxes = torch.cat([xs - width1, ys - height1, xs + width1, ys + height1],
dim=2)
detections = torch.cat([bboxes, scores, clses], dim=2)
return detections
def ctsegm_amodal_cmm_whiten_decode(hmap, regs, w_h_, codes_, offsets_, dictionary, code_range, K=100):
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap)
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
codes_ = codes_[0:1]
offsets_ = offsets_[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
codes_ = _tranpose_and_gather_feature(codes_, inds)
codes_ = codes_.view(batch, K, 64)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([xs - w_h_[..., 0:1] / 2,
ys - w_h_[..., 1:2] / 2,
xs + w_h_[..., 0:1] / 2,
ys + w_h_[..., 1:2] / 2], dim=2)
offsets_ = _tranpose_and_gather_feature(offsets_, inds)
# codes_ = codes_ * code_stat[1].view(1, 1, -1) + code_stat[0].view(1, 1, -1) # recover the original unnormalized codes
codes_ = (codes_ + 1) / 2. * (code_range[1] - code_range[0]) + code_range[0]
segms = torch.matmul(codes_, dictionary)
segms = segms.view(batch, K, 32, 2) + offsets_.view(batch, K, 1, 2) + \
torch.cat([xs, ys], dim=2).view(batch, K, 1, 2)
segmentations = torch.cat([segms.view(batch, K, -1), bboxes, scores, clses], dim=2)
return segmentations
def ctsegm_inmodal_norm_code_decode(hmap, regs, w_h_, codes_, offsets_, contour_std, dictionary, K=100):
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap)
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
codes_ = codes_[0:1]
offsets_ = offsets_[0:1]
contour_std = contour_std[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
contour_std = _tranpose_and_gather_feature(contour_std, inds)
contour_std = contour_std.view(batch, K, 1)
codes_ = _tranpose_and_gather_feature(codes_, inds)
codes_ = codes_.view(batch, K, 64)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([xs - w_h_[..., 0:1] / 2,
ys - w_h_[..., 1:2] / 2,
xs + w_h_[..., 0:1] / 2,
ys + w_h_[..., 1:2] / 2], dim=2)
offsets_ = _tranpose_and_gather_feature(offsets_, inds)
segms = torch.matmul(codes_, dictionary) * contour_std
segms = segms.view(batch, K, 32, 2) + offsets_.view(batch, K, 1, 2) + \
torch.cat([xs, ys], dim=2).view(batch, K, 1, 2)
segmentations = torch.cat([segms.view(batch, K, -1), bboxes, scores, clses], dim=2)
return segmentations
def ctsegm_fourier_decode(hmap, regs, w_h_, real_, imaginary_, K=100):
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap)
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
real_ = real_[0:1]
imaginary_ = imaginary_[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 4)
real_ = _tranpose_and_gather_feature(real_, inds)
real_ = real_.view(batch, K, 32, 1)
imaginary_ = _tranpose_and_gather_feature(imaginary_, inds)
imaginary_ = imaginary_.view(batch, K, 32, 1)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([xs - w_h_[..., 2:3],
ys - w_h_[..., 0:1],
xs + w_h_[..., 3:4],
ys + w_h_[..., 1:2]], dim=2)
complex_codes = torch.cat([real_, imaginary_], dim=3) * 32.
segms = torch.ifft(complex_codes, signal_ndim=1)
segms = segms + torch.cat([xs, ys], dim=2).view(batch, K, 1, 2)
segmentations = torch.cat([segms.view(batch, K, -1), bboxes, scores, clses], dim=2)
return segmentations
def ctsegm_shift_code_decode(hmap, regs, w_h_, codes_, dictionary, K=100):
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap)
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
codes_ = codes_[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 4)
codes_ = _tranpose_and_gather_feature(codes_, inds)
codes_ = codes_.view(batch, K, 64)
# codes_ = torch.log(codes_).view(batch, K, 64)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([xs - w_h_[..., 2:3],
ys - w_h_[..., 0:1],
xs + w_h_[..., 3:4],
ys + w_h_[..., 1:2]], dim=2)
segms = torch.matmul(codes_, dictionary)
segms = segms.view(batch, K, 32, 2) + torch.cat([xs, ys], dim=2).view(batch, K, 1, 2)
segmentations = torch.cat([segms.view(batch, K, -1), bboxes, scores, clses], dim=2)
return segmentations
def ctsegm_decode(hmap, regs, w_h_, codes_, dictionary, K=100):
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap)
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
codes_ = codes_[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
std_ = torch.sqrt(torch.sum(w_h_ ** 2., dim=2, keepdim=True))
codes_ = _tranpose_and_gather_feature(codes_, inds)
# codes_ = codes_.view(batch, K, 64)
codes_ = torch.log(codes_).view(batch, K, 64)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
segms = torch.matmul(codes_, dictionary)
# print('Sizes:', segms.size(), std_.size(), xs.size())
segms = (segms * std_).view(batch, K, 32, 2) + torch.cat([xs, ys], dim=2).view(batch, K, 1, 2)
segmentations = torch.cat([segms.view(batch, K, -1), scores, clses], dim=2)
return segmentations
def ctsegm_code_shape_decode(hmap, regs, w_h_, shapes_, K=100):
batch, cat, height, width = hmap.shape
hmap = torch.sigmoid(hmap)
# if flip test
if batch > 1:
hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
regs = regs[0:1]
shapes_ = shapes_[0:1]
batch = 1
hmap = _nms(hmap) # perform nms on heatmaps
scores, inds, clses, ys, xs = _topk(hmap, K=K)
regs = _tranpose_and_gather_feature(regs, inds)
regs = regs.view(batch, K, 2)
xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
ys = ys.view(batch, K, 1) + regs[:, :, 1:2]
w_h_ = _tranpose_and_gather_feature(w_h_, inds)
w_h_ = w_h_.view(batch, K, 2)
shapes_ = shapes_.view(batch, K, 64)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([xs - w_h_[..., 0:1] / 2,
ys - w_h_[..., 1:2] / 2,
xs + w_h_[..., 0:1] / 2,
ys + w_h_[..., 1:2] / 2], dim=2)
segmentations = torch.cat([shapes_, bboxes, scores, clses], dim=2)
return segmentations
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
# torch.autograd.set_detect_anomaly(mode=True)
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(
cfg.device, non_blocking=True)
dict_tensor.requires_grad = False
outputs = model(batch['image'])
# hmap, regs, w_h_, codes_1, codes_2, codes_3, offsets = zip(*outputs)
hmap, regs, w_h_, codes, offsets = zip(*outputs)
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
codes = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes
]
# c_2 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_2]
# c_3 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_3]
offsets = [
_tranpose_and_gather_feature(r, batch['inds']) for r in offsets
]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
offsets_loss = _reg_loss(offsets, batch['offsets'],
batch['ind_masks'])
# codes_loss = (norm_reg_loss(c_1, batch['codes'], batch['ind_masks'], sparsity=0.)
# + norm_reg_loss(c_2, batch['codes'], batch['ind_masks'], sparsity=0.)
# + norm_reg_loss(c_3, batch['codes'], batch['ind_masks'], sparsity=0.)) / 3.
if cfg.code_loss == 'norm':
codes_loss = norm_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
sparsity=0.)
elif cfg.code_loss == 'adapt':
codes_loss = adapt_norm_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
sparsity=0.,
norm=cfg.adapt_norm)
elif cfg.code_loss == 'wing':
codes_loss = wing_norm_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
sparsity=0.,
epsilon=cfg.wing_epsilon,
omega=cfg.wing_omega)
else:
print('Loss type for code not implemented yet.')
raise NotImplementedError
loss = 1. * hmap_loss + 1. * reg_loss + 0.1 * w_h_loss + 0.1 * offsets_loss + \
cfg.code_loss_weight * codes_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
'Loss: hmap = %.3f reg = %.3f w_h = %.3f code = %.3f offsets = %.3f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
codes_loss.item(), offsets_loss.item()) +
' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('offset_loss', offsets_loss.item(),
step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
return
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(
cfg.device, non_blocking=True)
dict_tensor.requires_grad = False
outputs = model(batch['image'])
hmap, regs, w_h_, codes_1, codes_2, codes_3, offsets = zip(
*outputs)
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
c_1 = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes_1
]
c_2 = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes_2
]
c_3 = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes_3
]
offsets = [
_tranpose_and_gather_feature(r, batch['inds']) for r in offsets
]
# shapes_1 = [torch.matmul(c, dict_tensor) for c in c_1]
# shapes_2 = [torch.matmul(c, dict_tensor) for c in c_2]
# shapes_3 = [torch.matmul(c, dict_tensor) for c in c_3]
hmap_loss = _neg_loss(hmap, batch['hmap'])
# occ_loss = _neg_loss(occ_map, batch['occ_map'], ex=4.0)
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
offsets_loss = _reg_loss(offsets, batch['offsets'],
batch['ind_masks'])
codes_loss = (
norm_reg_loss(c_1, batch['codes'], batch['ind_masks']) +
norm_reg_loss(c_2, batch['codes'], batch['ind_masks']) +
norm_reg_loss(c_3, batch['codes'], batch['ind_masks'])) / 3.
# cmm_loss = (contour_mapping_loss(c_1, shapes_1, batch['shapes'], batch['ind_masks'], roll=False)
# + contour_mapping_loss(c_2, shapes_2, batch['shapes'], batch['ind_masks'], roll=False)
# + contour_mapping_loss(c_3, shapes_3, batch['shapes'], batch['ind_masks'], roll=False)) / 3.
# cmm_loss = (_reg_loss(shapes_1, batch['shapes'], batch['ind_masks'])
# + _reg_loss(shapes_2, batch['shapes'], batch['ind_masks'])
# + _reg_loss(shapes_3, batch['shapes'], batch['ind_masks'])) / 3.
# loss = 1. * hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.cmm_loss_weight * cmm_loss \
# + cfg.code_loss_weight * codes_loss + 0.1 * offsets_loss
loss = 1 * hmap_loss + 1 * reg_loss + 0.1 * w_h_loss \
+ cfg.code_loss_weight * codes_loss + 0.1 * offsets_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
'Loss: hmap = %.3f reg = %.3f w_h = %.3f code = %.3f offsets = %.3f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
codes_loss.item(), offsets_loss.item()) +
' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
# summary_writer.add_scalar('occ_loss', occ_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('offset_loss', offsets_loss.item(),
step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
# summary_writer.add_scalar('cmm_loss', cmm_loss.item(), step)
return
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device, non_blocking=True)
dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device, non_blocking=True)
dict_tensor.requires_grad = False
outputs = model(batch['image'], batch['inds'], batch['centers'])
hmap, regs, w_h_, offsets, codes_1, codes_2, codes_3, shapes, shapes_1 = zip(*outputs)
regs = [_tranpose_and_gather_feature(r, batch['inds']) for r in regs]
w_h_ = [_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_]
c_1 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_1]
c_2 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_2]
c_3 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_3]
offsets = [_tranpose_and_gather_feature(r, batch['inds']) for r in offsets]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
offsets_loss = _reg_loss(offsets, batch['offsets'], batch['ind_masks'])
codes_loss = (sparse_reg_loss(c_1, batch['codes'], batch['ind_masks'])
+ sparse_reg_loss(c_2, batch['codes'], batch['ind_masks'])
+ sparse_reg_loss(c_3, batch['codes'], batch['ind_masks'])) / 3.
# shapes_loss = contour_mapping_loss(codes_, shapes_, batch['shapes'], batch['ind_masks'], roll=False)
mask = batch['ind_masks'][:, :, None].expand_as(batch['centered_shapes']).float()
shapes_loss = (sum(nn.functional.l1_loss(s * mask, batch['centered_shapes'] * mask, reduction='sum') / (mask.sum() + 1e-4) for s in shapes) \
+ sum(nn.functional.l1_loss(s * mask, batch['centered_shapes'] * mask, reduction='sum') / (mask.sum() + 1e-4) for s in shapes_1)) / len(shapes) / 2. * cfg.shape_loss_weight
# shapes_loss = (sum([nn.functional.l1_loss(p.view(p.size(0), p.size(1), -1) * mask, batch['shapes'] * mask, reduction='sum') / (mask.sum() + 1e-4) for p in polys_1])
# + sum([nn.functional.l1_loss(p.view(p.size(0), p.size(1), -1) * mask, batch['shapes'] * mask, reduction='sum') / (mask.sum() + 1e-4) for p in polys_2])
# + sum([nn.functional.l1_loss(p.view(p.size(0), p.size(1), -1) * mask, batch['shapes'] * mask, reduction='sum') / (mask.sum() + 1e-4) for p in shapes_])) / 3.
loss = 1 * hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.code_loss_weight * codes_loss + \
shapes_loss + 0.1 * offsets_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log('[%d/%d-%d/%d] ' % (epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss = %.3f reg_loss = %.3f w_h_loss = %.3f offsets = %.3f code_loss = %.3f shapes_loss = %.3f' %
(hmap_loss.item(), reg_loss.item(), w_h_loss.item(), offsets_loss.item(), codes_loss.item(), shapes_loss.item()) +
' (%d samples/sec)' % (cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
summary_writer.add_scalar('shape_loss', shapes_loss.item(), step)
summary_writer.add_scalar('offsets_loss', offsets_loss.item(), step)
return
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(
cfg.device, non_blocking=True)
dict_tensor.requires_grad = False
outputs = model(batch['image'])
hmap, regs, w_h_, offsets, codes_1, codes_2, codes_3, active_cls = zip(
*outputs)
# hmap, regs, w_h_, offsets, active_codes, inactive_codes, active_cls = zip(*outputs)
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
active_cls = [
_tranpose_and_gather_feature(r, batch['inds'])
for r in active_cls
]
offsets = [
_tranpose_and_gather_feature(r, batch['inds']) for r in offsets
]
active_mask = batch['active']
c_1 = [
_tranpose_and_gather_feature(r, batch['inds']) * active_mask
for r in codes_1
]
c_2 = [
_tranpose_and_gather_feature(r, batch['inds']) * active_mask
for r in codes_2
]
c_3 = [
_tranpose_and_gather_feature(r, batch['inds']) * active_mask
for r in codes_3
]
# active_codes = [_tranpose_and_gather_feature(r, batch['inds']) * active_mask for r in active_codes]
mask = batch['ind_masks'][:, :,
None].expand_as(batch['active']).float()
active_cls_loss = sum(
BCE(r * mask, batch['active'] * mask) / (mask.sum() + 1e-4)
for r in active_cls) / len(active_cls)
# active_cls_loss = _bce_loss(active_cls, batch['active'], batch['ind_masks'])
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
offsets_loss = _reg_loss(offsets, batch['offsets'],
batch['ind_masks'])
codes_loss = (sparse_reg_loss(c_1,
batch['codes'] * active_mask,
batch['ind_masks'],
sparsity=0) +
sparse_reg_loss(c_2,
batch['codes'] * active_mask,
batch['ind_masks'],
sparsity=0) +
sparse_reg_loss(c_3,
batch['codes'] * active_mask,
batch['ind_masks'],
sparsity=0)) / 3.
# codes_loss = norm_reg_loss(active_codes, batch['codes'] * active_mask, batch['ind_masks']) \
# + norm_reg_loss(inactive_codes, batch['codes'] * inactive_mask, batch['ind_masks'])
# cmm_loss = (contour_mapping_loss(c_1, shapes_1, batch['shapes'], batch['ind_masks'], roll=False)
# + contour_mapping_loss(c_2, shapes_2, batch['shapes'], batch['ind_masks'], roll=False)
# + contour_mapping_loss(c_3, shapes_3, batch['shapes'], batch['ind_masks'], roll=False)) / 3.
# cmm_loss = (_reg_loss(shapes_1, batch['shapes'], batch['ind_masks'])
# + _reg_loss(shapes_2, batch['shapes'], batch['ind_masks'])
# + _reg_loss(shapes_3, batch['shapes'], batch['ind_masks'])) / 3.
loss = 1 * hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.code_loss_weight * codes_loss \
+ 0.1 * offsets_loss + cfg.bce_loss_weight * active_cls_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
'Loss: hmap = %.3f reg = %.3f w_h = %.3f code = %.3f offsets = %.3f active = %.3f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
codes_loss.item(), offsets_loss.item(),
active_cls_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
# summary_writer.add_scalar('occ_loss', occ_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('offset_loss', offsets_loss.item(),
step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
summary_writer.add_scalar('active_cls_loss',
active_cls_loss.item(), step)
# summary_writer.add_scalar('cmm_loss', cmm_loss.item(), step)
return
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(
cfg.device, non_blocking=True)
dict_tensor.requires_grad = False
outputs = model(batch['image'])
hmap, regs, w_h_, codes, offsets, votes = zip(*outputs)
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
codes = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes
]
votes = [
_tranpose_and_gather_feature(r, batch['inds']) for r in votes
]
offsets = [
_tranpose_and_gather_feature(r, batch['inds']) for r in offsets
]
shapes = [torch.matmul(c, dict_tensor) for c in codes]
hmap_loss = _neg_loss(hmap, batch['hmap'])
bce_loss = _bce_loss(votes, batch['votes'], batch['ind_masks'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
offsets_loss = _reg_loss(offsets, batch['offsets'],
batch['ind_masks'])
if cfg.code_loss == 'norm':
codes_loss = norm_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
sparsity=0.01)
elif cfg.code_loss == 'adapt':
codes_loss = adapt_norm_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
norm=cfg.adapt_norm)
elif cfg.code_loss == 'wing':
codes_loss = wing_norm_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
sparsity=0.01,
epsilon=cfg.wing_epsilon,
omega=cfg.wing_omega)
elif cfg.code_loss == 'mse':
codes_loss = mse_reg_loss(codes,
batch['codes'],
batch['ind_masks'],
sparsity=0.001)
else:
print('Loss type for code not implemented yet.')
raise NotImplementedError
if cfg.shape_loss == 'cmm':
shape_loss = contour_mapping_loss(codes,
shapes,
batch['shapes'],
batch['ind_masks'],
roll=False)
elif cfg.shape_loss == 'piou':
shape_loss = 0.
for i in range(len(shapes)):
shape_loss += PIoU_loss(pred_shapes=shapes[i],
gt_shapes=batch['shapes'],
mask=batch['ind_masks'])
else:
print('Loss type for shape not implemented yet.')
raise NotImplementedError
loss = 1 * hmap_loss + 1 * reg_loss + 0.1 * w_h_loss + cfg.code_loss_weight * codes_loss \
+ 0.5 * offsets_loss + cfg.shape_loss_weight * shape_loss + cfg.bce_loss_weight * bce_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
'Loss: hmap = %.3f reg = %.3f w_h = %.3f code = %.3f offsets = %.3f shape = %.3f vote = %.3f'
%
(hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
codes_loss.item(), offsets_loss.item(), shape_loss.item(),
bce_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('offset_loss', offsets_loss.item(),
step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
summary_writer.add_scalar('vote_loss', bce_loss.item(), step)
summary_writer.add_scalar('shape_loss', shape_loss.item(),
step)
return
def train(epoch):
print_log('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
# dict_tensor = torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device, non_blocking=True)
# dict_tensor.requires_grad = False
outputs = model(batch['image'])
hmap, regs, w_h_, codes_1, offsets, votes = zip(*outputs)
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
c_1 = [
_tranpose_and_gather_feature(r, batch['inds']) for r in codes_1
]
# c_2 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_2]
# c_3 = [_tranpose_and_gather_feature(r, batch['inds']) for r in codes_3]
offsets = [
_tranpose_and_gather_feature(r, batch['inds']) for r in offsets
]
votes = [
_tranpose_and_gather_feature(r, batch['inds']) for r in votes
]
hmap_loss = _neg_loss(hmap, batch['hmap'])
vote_loss = _bce_loss(votes, batch['votes'], batch['ind_masks'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
offsets_loss = _reg_loss(offsets, batch['offsets'],
batch['ind_masks'])
# codes_loss = (norm_reg_loss(c_1, batch['codes'], batch['ind_masks'], sparsity=0.)
# + norm_reg_loss(c_2, batch['codes'], batch['ind_masks'], sparsity=0.)
# + norm_reg_loss(c_3, batch['codes'], batch['ind_masks'], sparsity=0.)) / 3.
codes_loss = norm_reg_loss(c_1,
batch['codes'],
batch['ind_masks'],
sparsity=0.0)
loss = 1. * hmap_loss + 1. * reg_loss + 0.1 * w_h_loss + 0.1 * offsets_loss + \
cfg.vote_loss_weight * vote_loss + cfg.code_loss_weight * codes_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print_log(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
'Loss: hmap = %.3f reg = %.3f w_h = %.3f code = %.3f offsets = %.3f votes = %.3f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
codes_loss.item(), offsets_loss.item(),
vote_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('vote_loss', vote_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('offset_loss', offsets_loss.item(),
step)
summary_writer.add_scalar('code_loss', codes_loss.item(), step)
return
batch[k] = batch[k].to(device=cfg.device, non_blocking=True)
'''
#TUDO
由于用的是CenterTrack的nuscenes数据集，则根据CenterTrack的代码格式继续写训练代码
数据集 √
模型√
loss
train
'''
outputs = model(batch['image'])
#hmap, regs, w_h_ = zip(*outputs)
hmap = outputs['hm']
regs = outputs['reg']
w_h_ = outputs['wh']
regs = [_tranpose_and_gather_feature(r, batch['inds']) for r in regs]
w_h_ = [_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_]
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _reg_loss(regs, batch['regs'], batch['ind_masks'])
w_h_loss = _reg_loss(w_h_, batch['w_h_'], batch['ind_masks'])
loss = hmap_loss + 1 * reg_loss + 0.1 * w_h_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print('[%d/%d-%d/%d] ' % (epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
