Pytorch：使用自定义权重图应用交叉熵损失

Question

我正在使用 pytorch 中的 u-net 架构解决多类分割问题。 正如U-NET 论文中所述，我正在尝试实现自定义权重图来应对类不平衡。

以下是我要应用的操作 -

此外，我减少了batch_size=1，以便在将其传递给precompute_to_masks 函数时删除该维度。 我尝试了以下方法-

def precompute_for_image(masks):
    masks = masks.cpu()
    cls = masks.unique()
    res = torch.stack([torch.where(masks==cls_val, torch.tensor(1), torch.tensor(0)) for cls_val in cls])
    return res

def train(n_epochs, loaders, model, optimizer, criterion, use_cuda, save_path):

        ###################
        # train the model #
        ###################
        model.train()
        for batch_idx, (data, target) in enumerate(final_train_loader):
            # move to GPU
            if use_cuda:
                data, target = data.cuda(), target.cuda()
            optimizer.zero_grad()
            output = model(data)
            temp_target = precompute_for_image(target)
            w = weight_map(temp_target)
            loss = criterion(output,target)
            loss = w*loss
            loss.backward()
            optimizer.step()

            train_loss = train_loss + ((1 / (batch_idx + 1)) * (loss.data - train_loss))

    return model

其中 weight_map 是计算我从here 获得的权重掩码的函数 我面临的问题是，当我应用以下方法时，我得到了memory error。
我正在使用 61gb 内存和 Tesla V100 GPU。 我真的认为我以不正确的方式应用它。 怎么做？
我从训练循环中省略了非必要的细节。 下面是我的weight_map 函数：

from skimage.segmentation import find_boundaries

w0 = 10
sigma = 5

def make_weight_map(masks):
    """
    Generate the weight maps as specified in the UNet paper
    for a set of binary masks.

    Parameters
    ----------
    masks: array-like
        A 3D array of shape (n_masks, image_height, image_width),
        where each slice of the matrix along the 0th axis represents one binary mask.

    Returns
    -------
    array-like
        A 2D array of shape (image_height, image_width)

    """
    nrows, ncols = masks.shape[1:]
    masks = (masks > 0).astype(int)
    distMap = np.zeros((nrows * ncols, masks.shape[0]))
    X1, Y1 = np.meshgrid(np.arange(nrows), np.arange(ncols))
    X1, Y1 = np.c_[X1.ravel(), Y1.ravel()].T
    for i, mask in enumerate(masks):
        # find the boundary of each mask,
        # compute the distance of each pixel from this boundary
        bounds = find_boundaries(mask, mode='inner')
        X2, Y2 = np.nonzero(bounds)
        xSum = (X2.reshape(-1, 1) - X1.reshape(1, -1)) ** 2
        ySum = (Y2.reshape(-1, 1) - Y1.reshape(1, -1)) ** 2
        distMap[:, i] = np.sqrt(xSum + ySum).min(axis=0)
    ix = np.arange(distMap.shape[0])
    if distMap.shape[1] == 1:
        d1 = distMap.ravel()
        border_loss_map = w0 * np.exp((-1 * (d1) ** 2) / (2 * (sigma ** 2)))
    else:
        if distMap.shape[1] == 2:
            d1_ix, d2_ix = np.argpartition(distMap, 1, axis=1)[:, :2].T
        else:
            d1_ix, d2_ix = np.argpartition(distMap, 2, axis=1)[:, :2].T
        d1 = distMap[ix, d1_ix]
        d2 = distMap[ix, d2_ix]
        border_loss_map = w0 * np.exp((-1 * (d1 + d2) ** 2) / (2 * (sigma ** 2)))
    xBLoss = np.zeros((nrows, ncols))
    xBLoss[X1, Y1] = border_loss_map
    # class weight map
    loss = np.zeros((nrows, ncols))
    w_1 = 1 - masks.sum() / loss.size
    w_0 = 1 - w_1
    loss[masks.sum(0) == 1] = w_1
    loss[masks.sum(0) == 0] = w_0
    ZZ = xBLoss + loss
    return ZZ

错误的追溯-

MemoryError                               Traceback (most recent call last)
<ipython-input-30-f0a595b8de7e> in <module>
      1 # train the model
      2 model_scratch = train(20, final_train_loader, unet, optimizer, 
----> 3                       criterion, train_on_gpu, 'model_scratch.pt')

<ipython-input-29-b481b4f3120e> in train(n_epochs, loaders, model, optimizer, criterion, use_cuda, save_path)
     24             loss = criterion(output,target)
     25             target.requires_grad = False
---> 26             w = make_weight_map(target)
     27             loss = W*loss
     28             loss.backward()

<ipython-input-5-e75a6281476f> in make_weight_map(masks)
     33         X2, Y2 = np.nonzero(bounds)
     34         xSum = (X2.reshape(-1, 1) - X1.reshape(1, -1)) ** 2
---> 35         ySum = (Y2.reshape(-1, 1) - Y1.reshape(1, -1)) ** 2
     36         distMap[:, i] = np.sqrt(xSum + ySum).min(axis=0)
     37     ix = np.arange(distMap.shape[0])

MemoryError:

Answer 1

您的final_train_loader 为您提供输入图像data 和预期的像素级标签target。我假设（遵循 pytorch 的约定）data 的形状为 B-3-H-W 和 dtype=torch.float。
更重要的是，target 的形状是 B-H-W 和 dtype=torch.long。

另一方面，make_weight_map 期望其输入是 C-H-W（C = 类数，而不是批量大小），类型为 numpy 数组。

尝试提供make_weight_map 输入掩码正如它所期望的那样，看看是否会遇到类似的错误。
我还建议您可视化生成的权重图 - 以确保您的函数执行您期望的操作。