forked from ccf-ai-infra/GPUCodeForces
77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
B, COORD = 32, 4
|
|
EPS = 1e-6
|
|
|
|
|
|
class IoULossBatch(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
px1, py1, px2, py2 = pred.split(1, dim=1)
|
|
tx1, ty1, tx2, ty2 = target.split(1, dim=1)
|
|
|
|
# 1. 计算交集区域坐标
|
|
ix1 = torch.max(px1, tx1)
|
|
iy1 = torch.max(py1, ty1)
|
|
ix2 = torch.min(px2, tx2)
|
|
iy2 = torch.min(py2, ty2)
|
|
|
|
# 2. 计算交集区域面积 (确保宽度/高度非负)
|
|
iw = torch.max(ix2 - ix1, torch.tensor(0.0).to(pred.device))
|
|
ih = torch.max(iy2 - iy1, torch.tensor(0.0).to(pred.device))
|
|
intersection = iw * ih
|
|
|
|
# 3. 计算预测框和目标框的面积
|
|
area_p = (px2 - px1) * (py2 - py1)
|
|
area_t = (tx2 - tx1) * (ty2 - ty1)
|
|
|
|
# 4. 计算并集区域面积: Union = Area_p + Area_t - Intersection
|
|
union = area_p + area_t - intersection
|
|
|
|
# 5. 计算 IoU
|
|
iou = intersection / (union + EPS)
|
|
|
|
# 6. 计算 IoU Loss (1 - IoU)
|
|
loss = 1.0 - iou
|
|
|
|
return loss.mean()
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.op = IoULossBatch()
|
|
|
|
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
return self.op(pred, target)
|
|
|
|
|
|
def get_inputs():
|
|
torch.manual_seed(42)
|
|
# 随机生成坐标 (0, 100) 范围的基数
|
|
base = torch.rand(B, COORD) * 100
|
|
|
|
# 构建预测框 (确保 x1 < x2 且 y1 < y2)
|
|
pred = torch.empty_like(base)
|
|
pred[:, 0] = torch.min(base[:, 0], base[:, 2])
|
|
pred[:, 1] = torch.min(base[:, 1], base[:, 3])
|
|
pred[:, 2] = torch.max(base[:, 0], base[:, 2])
|
|
pred[:, 3] = torch.max(base[:, 1], base[:, 3])
|
|
|
|
# 构建目标框
|
|
target = torch.empty_like(base)
|
|
target[:, 0] = torch.min(base[:, 0] + 5, base[:, 2] + 5)
|
|
target[:, 1] = torch.min(base[:, 1] + 5, base[:, 3] + 5)
|
|
target[:, 2] = torch.max(base[:, 0] + 5, base[:, 2] + 5)
|
|
target[:, 3] = torch.max(base[:, 1] + 5, base[:, 3] + 5)
|
|
|
|
return [pred, target]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |