背景:一个图像预测任务,图像并不是指常见的jpg等格式的图片,其实就是网格内每个像素点有数值的那种类似于热点图的图像,如下图
问题:如上图所示,一张图里其实有很大一片区域是0值。但搭的网络输入输出是整图,维度为(32,224,224)。最近意识到计算loss不可以将整图放入计算,只想要关注body内有值的区域,所以就用上自写的loss function,出现了计算极其缓慢的现象。
——
train预测过程:x为输入,y为GT,y_pred为预测结果,b是与y维度相同的、只有0或1的mask,b==1就是body内区域。我的做法是把b展平成一维,再用np.where读取b==1的坐标到BODY,同时y,y_pred也展平成一维。放入自写的smooth_L1计算loss
x = np.array(inputList)
y = np.array(labelList)
b = np.array(bodyList)
b = b.flatten()
BODY = np.where(b == 1)
BODY = np.array(BODY)
BODY = BODY.flatten()
x = np.float32(x)
y = np.float32(y)
x = torch.tensor(x)
y = torch.tensor(y)
BODY = torch.tensor(BODY)
x = x.to(device)
y = y.to(device)
BODY = BODY.to(device)
optimizer.zero_grad()
y_pred = model(x)
y_pred = y_pred.view(-1)
y = y.view(-1)
train_loss = smooth_l1(y_pred, y, BODY)
print(f"第 {t + 1}/{epoch} epoch, 第 {i + 1}/{batch_train} batch, train_loss: {train_loss.item()}")
train_loss.backward()
optimizer.step()
自写的loss:就是一个smoothL1 Loss,但是我只想计算body内区域。所以用了一个for循环读取mask里(也就是BODY里)存好的坐标。
# 定义损失函数
def smooth_item(x, beta):
if x < beta:
loss_item = (0.5*torch.square(x))/beta
else:
loss_item = x-0.5*beta
return loss_item
def smooth_l1(y_pred, y, mask, beta=1):
n = len(mask)
loss = 0
for i in mask:
loss += smooth_item(torch.abs(y_pred[i]-y[i]), beta)
return loss/n
debug发现这样子计算loss好慢好慢,有什么方法优化吗?要说用GPU加速,但这样子一个个坐标读取再累加也不是一个并行过程呀