pytorch模型中有单个参数需要学习

import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.model1 = nn.Conv2d(3, 20, kernel_size=3, stride=1)
        self.model2 = nn.Conv2d(3, 20, kernel_size=3, stride=1)
        self.a = nn.Parameter(torch.FloatTensor([0.1]))
        self.b = nn.Parameter(torch.FloatTensor([0.9]))

    def forward(self, inputs):
        # f1 = self.model1(inputs)
        # f2 = self.model2(inputs)
        # # return self.a * f1 + self.b * f2
        # return f1 + f2
        pass


def get_parameter_number(net):
    total_num = sum(p.numel() for p in net.parameters())
    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}


if __name__ == '__main__':
    modul = Model().cuda()
    print(get_parameter_number(modul))

分别输出定义了ab的模型和没定义ab的模型参数，一个是1122，一个是1120，相差两个，说明已经将a和b定义为要学习的变量