Градиентный вывод через пользовательскую функцию потерь

Я совершенно нов в Pytorch (и в машинном обучении в целом), поэтому мне трудно понять, что происходит в отношении настраиваемой функции потерь, которую я рассматриваю. Я понимаю, что происходит в функции, но мне нужно понять, как рассчитывается градиентный вывод последнего слоя сети.

ПРИМЕЧАНИЕ: Моя задача – внедрить эту настраиваемую функцию потерь в нашу индивидуальную библиотеку машинного обучения на C++. Но для этого мне нужно “вручную” вычислить градиенты для последнего слоя сети.

Итак, если я запускаю свою сеть:

results = my_network(inputs)

А затем мою функцию потерь:

loss = my_loss_fn(inputs, results, targets)
loss.backward()

Наконец, если я выведу цепочку grad_fn для “results”, я увижу:

цепочка grad_fn "results":
SqueezeBackward1
DivBackward0
SliceBackward
SqueezeBackward1
Col2ImBackward
TransposeBackward0
MulBackward0
FftC2RBackward
ViewAsComplexBackward
TransposeBackward0
ViewBackward
LeakyReluBackward0 <=== Это последний слой my_network
NativeBatchNormBackward
SlowConvTranspose2DBackward

Так что, если я правильно понимаю автограду, мне нужно будет реализовать каждую из этих grad_fn, чтобы получить градиентный вывод для LeakReluLayer?

ИЗМЕНЕНИЕ: Для сети я просто использую пару простых слоев. Она не предназначена для выполнения какой-либо задачи:

class NetTest(nn.Module):
def __init__(self, in_chans, out_chans, kernel, stride, padding):
    super().__init__()

    #self.conv = nn.Conv2d(in_channels=in_chans, out_channels=out_chans, kernel_size=kernel, stride=stride, padding=padding)
    self.trans = nn.ConvTranspose2d(in_channels=in_chans, out_channels=out_chans, kernel_size=kernel, stride=stride, padding=padding, output_padding=0)
    self.batchnorm = nn.BatchNorm2d(num_features=out_chans, eps=0,momentum=.1)
    self.leaky_relu = nn.LeakyReLU()

def forward(self, x):
    #x = self.conv(x)
    x = self.trans(x)
    x = self.batchnorm(x)
    x = self.leaky_relu(x)

    return x

Функция потерь выглядит следующим образом (не уверен, откуда она взялась):

def wsdr_fn(x_, y_pred_, y_true_, eps=1e-8):
    # x_ - это вход сети в формате STFT
    # y_pred_ - это результаты сети. Они уже обработаны через istft
    # y_true_ - это цель в формате STFT
    y_true_ = torch.squeeze(y_true_, 1)
    y_true = torch.istft(y_true_, n_fft=N_FFT, hop_length=HOP_LENGTH, normalized=True)
    x_ = torch.squeeze(x_, 1)
    x = torch.istft(x_, n_fft=N_FFT, hop_length=HOP_LENGTH, normalized=True)
    y_pred = y_pred_.flatten(1)
    y_true = y_true.flatten(1)
    x = x.flatten(1)

    def sdr_fn(true, pred, eps=1e-8):
            num = torch.sum(true * pred, dim=1)
            den = torch.norm(true, p=2, dim=1) * torch.norm(pred, p=2, dim=1)
        return -(num / (den + eps))

    # истинный иEstimated шум
    z_true = x - y_true
    z_pred = x - y_pred

    a1 = torch.sum(y_true**2, dim=1)   
    a = torch.sum(y_true**2, dim=1) / (torch.sum(y_true**2, dim=1) + torch.sum(z_true**2, dim=1) + eps)
    wSDR = a * sdr_fn(y_true, y_pred) + (1 - a) * sdr_fn(z_true, z_pred)

    return torch.mean(wSDR)

Во-первых, довольно нестандартно передавать входные данные в функцию потерь, если только у вас нет настраиваемой функции потерь, которая требует их специально. Это может происходить, например, когда вы пытаетесь сделать распределение выводов соответствующим (или наоборот игнорировать) распределение, основанное на входах, но это не так уж часто. Можете объяснить, что делает ваша функция потерь?

Во-вторых, относительно вычисления градиента, последний слой обычно является самым простым для ручного вычисления градиентов. Прямолинейно вычислить частную производную потерь по параметрам конечного слоя, так как это связано только с дифференцированием двух функций. Так что, нет, вам не нужно заново реализовывать все. Можете поделиться графом вашей модели, или это конфиденциально?