PyTorch中如何進(jìn)行模型的參數(shù)初始化

小樊
192
2024-03-05 19:31:04

在PyTorch中，可以通過(guò)定義一個(gè)函數(shù)來(lái)對(duì)模型的參數(shù)進(jìn)行初始化。一般情況下，PyTorch提供了一些內(nèi)置的初始化方法，如torch.nn.init模塊中的一些函數(shù)。以下是一種常見(jiàn)的初始化方法：

import torch
import torch.nn as nn
import torch.nn.init as init

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(100, 10)

    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

model = MyModel()
model.initialize_weights()

在上面的代碼中，我們定義了一個(gè)MyModel類(lèi)，其中包含一個(gè)線(xiàn)性層nn.Linear(100, 10)。使用initialize_weights函數(shù)對(duì)模型的參數(shù)進(jìn)行初始化，其中我們使用了Xavier初始化方法對(duì)權(quán)重進(jìn)行初始化，并將偏置初始化為0。您也可以根據(jù)需要選擇其他初始化方法。

0 踩