Running out of memory when calculating loss using SigmoidFocalCrossEntropy

Code Versions:

python == 3.8
tensorflow == 2.2.0
tensorflow-addons == 0.11.2

Recently I’ve been using tensorflow addon’s focal loss function for one of my models. In order to better understand/demonstrate what’s going on, I’ve been trying to recreate some of the graphs from the original paper – specifically, figure 4, which shows a cumulative sum plot of the minority and majority loss separately using different gammas on a converged model.

def calcCSP(pred, Y_test, g, label):

    true = np.asarray(Y_test)
    idsPos = np.where(true == 1)
    idsNeg = np.where(true == 0)

    pPos = pred[idsPos]
    yPos = true[idsPos]

    pNeg = pred[idsNeg]
    yNeg = true[idsNeg]

    if label == 1:
        p = pPos
        y = yPos
        title = "Positive Loss Distribution"
    else:
        p = pNeg
        y = yNeg
        title = "Negative Loss Distribution"

    p = tf.convert_to_tensor(p)
    y = tf.cast(y, tf.float32)

    fl = tfa.losses.SigmoidFocalCrossEntropy(alpha=PARAMS['alpha'], gamma=g)

    loss = fl(y, p)

    x = np.sort(loss)

    # Normalized Data
    x = x/sum(x)

    cdf = np.cumsum(x)
    n = len(x)
    share_of_population = np.arange(1, n + 1) / n


    cdf_materials = {"shares": share_of_population,
                     "cusum": cdf}

    return cdf_materials


cdfListPos = []
cdfListNeg = []

gammas = [0, 2, 4, 6, 8]

for g in gammas:
    # will need to store all this
    cdf_matPos = calcCDF(preds, y_test, g, 1)
    cdf_matNeg = calcCDF(preds, y_test, g, 0)
    cdfListPos.append(cdf_matPos)
    cdfListNeg.append(cdf_matNeg)

posplot = plt.figure()
for i in range(len(gammas)):
    plt.plot(cdfListPos[i]['shares'], cdfListPos[i]['cusum'], label= r'$gamma$ = ' + str(gammas[i]))
    plt.title('Positive Points CSP')
    plt.ylabel('Cumulative Normalized Loss')
    plt.legend()

negplot = plt.figure()
for i in range(len(gammas)):
    plt.plot(cdfListNeg[i]['shares'], cdfListNeg[i]['cusum'], label=r'$gamma$ = ' + str(gammas[i]))
    plt.title('Negative Points CSP')
    plt.ylabel('Cumulative Normalized Loss')
    plt.legend()

I’ve managed to run my code okay on smaller datasets – however, once the dataset gets too large, I seem to run out of memory.

W tensorflow/core/common_runtime/bfc_allocator.cc:434] Allocator (mklcpu) ran out of memory trying to allocate 4.47GiB (rounded to 4804430848)
Current allocation summary follows.

I’ve omitted the summary for brevity, but can add it in if necessary. I’m unsure of how to get around this. I’ve tried using tf.split to divide up the predicted and true values into smaller tensors, then passing them to the loss function separately and concatenating, but testing with a smaller dataset shows that the results don’t match when I compare them to the full versions. I’m honestly stumped on how best to work around this and would be open to any recommendations.