PyTorch計(jì)算損失函數(shù)對(duì)模型參數(shù)的Hessian矩陣示例

更新時(shí)間：2022年05月10日 10:59:43 作者：Cyril_KI

這篇文章主要為大家介紹了PyTorch計(jì)算損失函數(shù)對(duì)模型參數(shù)的Hessian矩陣的示例詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

前言

在實(shí)現(xiàn)Per-FedAvg的代碼時(shí)，遇到如下問題：

在這里插入圖片描述

可以發(fā)現(xiàn)，我們需要求損失函數(shù)對(duì)模型參數(shù)的Hessian矩陣。

模型定義

我們定義一個(gè)比較簡(jiǎn)單的模型：

class ANN(nn.Module):
    def __init__(self):
        super(ANN, self).__init__()
        self.sigmoid = nn.Sigmoid()
        self.fc1 = nn.Linear(3, 4)
        self.fc2 = nn.Linear(4, 5)

    def forward(self, data):
        x = self.fc1(data)
        x = self.fc2(x)

        return x

輸出一下模型的參數(shù)：

model = ANN()
for param in model.parameters():
    print(param.size())

輸出如下：

torch.Size([4, 3])
torch.Size([4])
torch.Size([5, 4])
torch.Size([5])

求解Hessian矩陣

我們首先定義數(shù)據(jù)：

data = torch.tensor([1, 2, 3], dtype=torch.float)
label = torch.tensor([1, 1, 5, 7, 8], dtype=torch.float)
pred = model(data)
loss_fn = nn.MSELoss()
loss = loss_fn(pred, label)

然后求解一階梯度：

grads = torch.autograd.grad(loss, model.parameters(), retain_graph=True, create_graph=True)

輸出一下grads：

(tensor([[-1.0530, -2.1059, -3.1589],
        [ 2.3615,  4.7229,  7.0844],
        [-1.5046, -3.0093, -4.5139],
        [-2.0272, -4.0543, -6.0815]], grad_fn=<TBackward0>), tensor([-1.0530,  2.3615, -1.5046, -2.0272], grad_fn=<SqueezeBackward1>), tensor([[ 0.2945, -0.2725, -0.8159, -0.6720],
        [ 0.1936, -0.1791, -0.5362, -0.4416],
        [ 1.0800, -0.9993, -2.9918, -2.4641],
        [ 1.3448, -1.2444, -3.7255, -3.0683],
        [ 1.2436, -1.1507, -3.4450, -2.8373]], grad_fn=<TBackward0>), tensor([-0.6045, -0.3972, -2.2165, -2.7600, -2.5522],
       grad_fn=<MseLossBackwardBackward0>))

可以發(fā)現(xiàn)一共4個(gè)Tensor，分別為損失函數(shù)對(duì)四個(gè)參數(shù)Tensor（兩層，每層都有權(quán)重和偏置）的梯度。

然后針對(duì)每一個(gè)Tensor求解二階梯度：

hessian_params = []
    for k in range(len(grads)):
        hess_params = torch.zeros_like(grads[k])
        for i in range(grads[k].size(0)):
            # 判斷是w還是b
            if len(grads[k].size()) == 2:
                # w
                for j in range(grads[k].size(1)):
                    hess_params[i, j] = torch.autograd.grad(grads[k][i][j], model.parameters(), retain_graph=True)[k][i, j]
            else:
                # b
                hess_params[i] = torch.autograd.grad(grads[k][i], model.parameters(), retain_graph=True)[k][i]
        hessian_params.append(hess_params)

這里需要注意：由于模型一共兩層，每一層都有權(quán)重和偏置，其中權(quán)重參數(shù)為二維，偏置參數(shù)為一維，在進(jìn)行具體的二階梯度求導(dǎo)時(shí)，需要進(jìn)行判斷。

最終得到的hessian_params是一個(gè)列表，列表中包含四個(gè)Tensor，對(duì)應(yīng)損失函數(shù)對(duì)兩層網(wǎng)絡(luò)權(quán)重和偏置的二階梯度。

以上就是PyTorch計(jì)算損失函數(shù)對(duì)模型參數(shù)的Hessian矩陣示例的詳細(xì)內(nèi)容，更多關(guān)于PyTorch計(jì)算損失函數(shù)Hessian矩陣的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

您可能感興趣的文章: