Capturea.PNG1379×864 98.8 KB

class MyNet1(nn.Module):
def __init__(self):
    super(MyNet1, self).__init__()

    self.common = nn.Sequential(
        nn.Linear(64*64*3,32*32*3),
        nn.ReLU(),
        nn.Dropout(0.2),
        nn.Linear(32*32*3, 16*16*3),
        nn.ReLU(),
        nn.Dropout(0.2),
        nn.Linear(16*16*3,320),
        nn.Dropout(0.2) 
    )


    self.branch1 = nn.Sequential(
       nn.Linear(320,320),
       nn.ReLU(),
       nn.Dropout(0.2),
       nn.Linear(320,160),
       nn.ReLU(),
       nn.Dropout(0.2),
       nn.Linear(160,10)
       
    )

    self.branch2 =nn.Sequential(
       nn.Linear(320,320),
       nn.ReLU(),
       nn.Dropout(0.2),
       nn.Linear(320,160),
       nn.ReLU(),
       nn.Dropout(0.2),
       nn.Linear(160,10)          
    )


def forward(self, X):
        
        X = X.view(X.size(0),-1)
        
        X= self.common(X)
         

        op1= self.branch1(X)
        op2= self.branch2(X)
     
        
        return torch.cat([op1, op2], 1)