PyTorch多模型定义报错：实例化Model2触发TypeError问题求助

我想要在PyTorch中使用两个不同的模型，执行以下代码后无法成功运行第二个模型Model2，报错信息如下：

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(2, 64)
        self.linear2 = nn.Linear(64, 3)

    def forward(self, x):   
        x = self.linear1(x)
        x = torch.sigmoid(x)  
        x = self.linear2(x)

        return x
                
class Model2(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(2, 64)
        self.linear2 = nn.Linear(64, 1)

    def forward(self, x):   
        x = self.linear1(x)
        x = torch.sigmoid(x)  
        x = self.linear2(x)

        return x
                
net = Model()
net2 = Model2()

错误信息 #

TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_477/2280223066.py in <module>
     26 
     27 net = Model()
---&gt; 28 net2 = Model2()

/tmp/ipykernel_477/2280223066.py in __init__(self)
     14 class Model2(nn.Module):
     15     def __init__(self):
---&gt; 16         super(Model, self).__init__()
     17         self.linear1 = nn.Linear(2, 64)
     18         self.linear2 = nn.Linear(64, 1)

TypeError: super(type, obj): obj must be an instance or subtype of type

问题原因 #

Model2的__init__方法中，调用super()时传入的第一个参数错误，写成了Model，但这里必须传入当前类Model2。super()的第一个参数需要对应当前子类的类型，才能正确定位父类nn.Module的初始化方法。

修复后的代码 #

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(2, 64)
        self.linear2 = nn.Linear(64, 3)

    def forward(self, x):   
        x = self.linear1(x)
        x = torch.sigmoid(x)  
        x = self.linear2(x)

        return x
                
class Model2(nn.Module):
    def __init__(self):
        # 将第一个参数改为当前类Model2
        super(Model2, self).__init__()
        self.linear1 = nn.Linear(2, 64)
        self.linear2 = nn.Linear(64, 1)

    def forward(self, x):   
        x = self.linear1(x)
        x = torch.sigmoid(x)  
        x = self.linear2(x)

        return x
                
net = Model()
net2 = Model2()

更简洁的写法（Python3+） #

可以使用Python3的简化语法super().__init__()，无需手动传入类名和self，避免参数写错：

class Model2(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(2, 64)
        self.linear2 = nn.Linear(64, 1)

    def forward(self, x):   
        x = self.linear1(x)
        x = torch.sigmoid(x)  
        x = self.linear2(x)

        return x

内容的提问来源于stack exchange，提问作者asdfe