#Def the Network and forward prop

dim = 12288 #64*64*3channels
n_h = 7
n_classes = 2 # 2 classes
n_y = 5

class Net(nn.Module): 
    def __init__(self):
        
        super(Net,self).__init__()   
        self.input_layer = nn.Linear(dim,n_h)
        self.hidden_layer = nn.Linear(n_h,n_y) 
        self.output_layer = nn.Linear(n_y,n_classes) 
           
    def forward(self, x):
        x = x.view(-1, dim) #flatten image input
        x = F.relu(self.input_layer(x)) 
        x = F.relu(self.hidden_layer(x))
        x = self.output_layer(x)#no ReLU to the last layer.
        return F.log_softmax(x, dim=1) #convert logits to probabilities that sum up to 1. 

net= Net()
print(net)
summary(Net(), input_size=(1,dim), batch_size=bs, device="cpu") 
net.state_dict()

type or paste code here

Net(
(input_layer): Linear(in_features=12288, out_features=7, bias=True)
(hidden_layer): Linear(in_features=7, out_features=5, bias=True)
(output_layer): Linear(in_features=5, out_features=2, bias=True)
)
    Layer (type)               Output Shape         Param #

================================================================
Linear-1                    [32, 7]          86,023
Linear-2                    [32, 5]              40
Linear-3                    [32, 2]              12
Total params: 86,075
Trainable params: 86,075
Non-trainable params: 0
Input size (MB): 1.50
Forward/backward pass size (MB): 0.00
Params size (MB): 0.33
Estimated Total Size (MB): 1.83
This is the output when feeding hdf5 dataset with custom dataset class.
Instead, for MNIST:
Net(
(input_layer): Linear(in_features=784, out_features=7, bias=True)
(hidden_layer): Linear(in_features=7, out_features=5, bias=True)
(output_layer): Linear(in_features=5, out_features=10, bias=True)
)
    Layer (type)               Output Shape         Param #

================================================================
Linear-1                [-1, 32, 7]           5,495
Linear-2                [-1, 32, 5]              40
Linear-3               [-1, 32, 10]              60
Total params: 5,595
Trainable params: 5,595
Non-trainable params: 0
Input size (MB): 0.10
Forward/backward pass size (MB): 0.01
Params size (MB): 0.02
Estimated Total Size (MB): 0.12