pytorchで学習済みモデルをロードした後のfc層の重みを変更したい

Question

学習したモデルをロードした後に、fc層の重みをすべて0に変更したいのですが、上手くいかなくて困っています。
以下のコードでfc層の重みを0に変更しようと試みています。
しかし、出力結果がfc層の重みを0に変更する前と同じ結果が出力されているため、重みの変更に失敗していることが考えられます。

model = torch.nn.DataParallel(model).cuda()
model_path = 'model.pth'
model.load_state_dict(torch.load(model_path))

# fc.weight.shape = (10, 2048)
model.state_dict()['fc.weight'][0:10] = torch.ones(10,2048)

解決方法をご存じの方がいましたら、ご教授頂けると幸いです。
よろしくお願いいたします。

Answer 1

以下のようにすれば、モデルの FC 層の重みをすべて0にできます。

Pytorch - 重みの初期化手法と各モジュールのデフォルトの初期化方法

import torch.nn as nn

for m in model.modules():
    if isinstance(m, nn.Linear):
        nn.init.zeros_(m.weight)
        nn.init.zeros_(m.bias)