43.pdf

    1. import  torch
    2. from    torch import nn
    3. from    torch import optim
    7. class MyLinear(nn.Module):
    9.     def __init__(self, inp, outp):
    10.         super(MyLinear, self).__init__()
    12.         # requires_grad = True
    13.         self.w = nn.Parameter(torch.randn(outp, inp))
    14.         self.b = nn.Parameter(torch.randn(outp))
    16.     def forward(self, x):
    17.         x = x @ self.w.t() + self.b
    18.         return x
    21. class Flatten(nn.Module):
    23.     def __init__(self):
    24.         super(Flatten, self).__init__()
    26.     def forward(self, input):
    27.         return input.view(input.size(0), -1)
    31. class TestNet(nn.Module):
    33.     def __init__(self):
    34.         super(TestNet, self).__init__()
    36.         self.net = nn.Sequential(nn.Conv2d(1, 16, stride=1, padding=1),
    37.                                  nn.MaxPool2d(2, 2),
    38.                                  Flatten(),
    39.                                  nn.Linear(1*14*14, 10))
    41.     def forward(self, x):
    42.         return self.net(x)
    45. class BasicNet(nn.Module):
    47.     def __init__(self):
    48.         super(BasicNet, self).__init__()
    50.         self.net = nn.Linear(4, 3)
    52.     def forward(self, x):
    53.         return self.net(x)
    57. class Net(nn.Module):
    59.     def __init__(self):
    60.         super(Net, self).__init__()
    62.         self.net = nn.Sequential(BasicNet(),
    63.                                  nn.ReLU(),
    64.                                  nn.Linear(3, 2))
    66.     def forward(self, x):
    67.         return self.net(x)
    73. def main():
    74.     device = torch.device('cuda')
    75.     net = Net()
    76.     net.to(device)
    78.     net.train()
    80.     net.eval()
    82.     # net.load_state_dict(torch.load('ckpt.mdl'))
    83.     #
    84.     #
    85.     # torch.save(net.state_dict(), 'ckpt.mdl')
    87.     for name, t in net.named_parameters():
    88.         print('parameters:', name, t.shape)
    90.     for name, m in net.named_children():
    91.         print('children:', name, m)
    94.     for name, m in net.named_modules():
    95.         print('modules:', name, m)
    104. if __name__ == '__main__':
    105.     main()