forked from ccf-ai-infra/GPUCodeForces
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask):
|
|
pred_hm = torch.clamp(pred_hm, 1e-6, 1 - 1e-6)
|
|
|
|
pos_inds = gt_hm.eq(1).float()
|
|
neg_inds = gt_hm.lt(1).float()
|
|
neg_weights = torch.pow(1 - gt_hm, 4)
|
|
|
|
pos_loss = torch.log(pred_hm) * torch.pow(1 - pred_hm, 2) * pos_inds
|
|
neg_loss = torch.log(1 - pred_hm) * torch.pow(pred_hm, 2) * neg_weights * neg_inds
|
|
|
|
num_pos = pos_inds.sum()
|
|
pos_loss_sum = pos_loss.sum()
|
|
neg_loss_sum = neg_loss.sum()
|
|
|
|
if num_pos > 0:
|
|
hm_loss = -(pos_loss_sum + neg_loss_sum) / num_pos
|
|
else:
|
|
hm_loss = -neg_loss_sum
|
|
|
|
mask_expanded = mask.expand_as(pred_wh)
|
|
wh_loss = torch.sum(torch.abs(pred_wh - gt_wh) * mask_expanded)
|
|
reg_loss = torch.sum(torch.abs(pred_reg - gt_reg) * mask_expanded)
|
|
|
|
if num_pos > 0:
|
|
wh_loss = wh_loss / num_pos
|
|
reg_loss = reg_loss / num_pos
|
|
|
|
return hm_loss + 0.1 * wh_loss + 1.0 * reg_loss
|
|
|
|
|
|
batch_size = 4
|
|
channels = 4
|
|
height = 128
|
|
width = 128
|
|
|
|
|
|
def get_inputs():
|
|
pred_hm = torch.sigmoid(torch.randn(batch_size, channels, height, width))
|
|
gt_hm = torch.bernoulli(torch.full((batch_size, channels, height, width), 0.1))
|
|
pred_wh = torch.randn(batch_size, 2, height, width)
|
|
gt_wh = torch.randn(batch_size, 2, height, width)
|
|
pred_reg = torch.randn(batch_size, 2, height, width)
|
|
gt_reg = torch.randn(batch_size, 2, height, width)
|
|
mask = torch.bernoulli(torch.full((batch_size, 1, height, width), 0.1))
|
|
return [pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |