查看模型各层的输入输出维度

Handsome_zhou

模型各层的输入输出维度在搭建好模型之后，进行模型实例化，就可以打印出来了。

1. device = "cuda" if torch.cuda.is_available() else "cpu"
   
2. print(f"Using {device} device")
   
3. 
   
4. 
   
5. class NeuralNetwork(nn.Module):
   
6.     def __init__(self):
   
7.         super(NeuralNetwork, self).__init__()
   
8.         self.flatten = nn.Flatten()
   
9.         self.linear_relu_stack = nn.Sequential(
   
10.             nn.Linear(28*28, 512),
    
11.             nn.ReLU(),
    
12.             nn.Linear(512, 512),
    
13.             nn.ReLU(),
    
14.             nn.Linear(512,10)
    
15.         )
    
16.    
    
17.     def forward(self, x):
    
18.         x = self.flatten(x)
    
19.         logits = self.linear_relu_stack(x)
    
20.         return logits
    
21. 
    
22. model = NeuralNetwork().to(device)
    
23. print(model)

复制代码

输出结果:

1. import json
   
2. import multiprocessing #使用python的多线程
   
3. import os
   
4. import torch
   
5. from torch import nn
   
6. from d2l import torch as d2l
   
7. 
   
8. #load bert.small预训练模型
   
9. devices = d2l.try_all_gpus()
   
10. bert,vocab = load_pretrained_model('bert.small',num_hiddens=256,ffn_num_hiddens=512,
    
11.                                  num_heads=4,num_layers=2,dropout=0.1,max_len=512,devices=devices)
    
12. 
    
13. class BERTClassifier(nn.Module):
    
14.     def __init__(self, bert):
    
15.         super(BERTClassifier, self).__init__()
    
16.         self.encoder = bert.encoder
    
17.         self.hidden = bert.hidden
    
18.         self.output = nn.Linear(256,3)
    
19.         
    
20.     def forward(self, inputs):
    
21.         tokens_X, segments_X, valid_lens_x = inputs
    
22.         encoded_X = self.encoder(tokens_X, segments_X, valid_lens_x)
    
23.         return self.output(self.hidden(encoded_X[:, 0, :]))
    
24.    
    
25. net = BERTClassifier(bert).to(device)#模型实例化
    
26. print(net)

复制代码

输出结果:



textCNN模型的各层输入输出维度：

1. import torch
   
2. import torch.nn as nn
   
3. from torch.nn import functional as F
   
4. import math
   
5. 
   
6. class textCNN(nn.Module):
   
7.     def __init__(self, param):
   
8.         super(textCNN, self).__init__()
   
9.         ci = 1  # input chanel size
   
10.         kernel_num = param['kernel_num'] # output chanel size
    
11.         kernel_size = param['kernel_size']
    
12.         vocab_size = param['vocab_size']
    
13.         embed_dim = param['embed_dim']
    
14.         dropout = param['dropout']
    
15.         class_num = param['class_num']
    
16.         self.param = param
    
17.         self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=1)
    
18.         self.conv11 = nn.Conv2d(ci, kernel_num, (kernel_size[0], embed_dim))
    
19.         self.conv12 = nn.Conv2d(ci, kernel_num, (kernel_size[1], embed_dim))
    
20.         self.conv13 = nn.Conv2d(ci, kernel_num, (kernel_size[2], embed_dim))
    
21.         self.dropout = nn.Dropout(dropout)
    
22.         self.fc1 = nn.Linear(len(kernel_size) * kernel_num, class_num)
    
23. 
    
24.     def init_embed(self, embed_matrix):
    
25.         self.embed.weight = nn.Parameter(torch.Tensor(embed_matrix))
    
26. 
    
27.     @staticmethod
    
28.     def conv_and_pool(x, conv):
    
29.         # x: (batch, 1, sentence_length,  )
    
30.         x = conv(x)
    
31.         # x: (batch, kernel_num, H_out, 1)
    
32.         x = F.relu(x.squeeze(3))
    
33.         # x: (batch, kernel_num, H_out)
    
34.         x = F.max_pool1d(x, x.size(2)).squeeze(2)
    
35.         #  (batch, kernel_num)
    
36.         return x
    
37. 
    
38.     def forward(self, x):
    
39.         # x: (batch, sentence_length)
    
40.         x = self.embed(x)
    
41.         # x: (batch, sentence_length, embed_dim)
    
42.         # TODO init embed matrix with pre-trained
    
43.         x = x.unsqueeze(1)
    
44.         # x: (batch, 1, sentence_length, embed_dim)
    
45.         x1 = self.conv_and_pool(x, self.conv11)  # (batch, kernel_num)
    
46.         x2 = self.conv_and_pool(x, self.conv12)  # (batch, kernel_num)
    
47.         x3 = self.conv_and_pool(x, self.conv13)  # (batch, kernel_num)
    
48.         x = torch.cat((x1, x2, x3), 1)  # (batch, 3 * kernel_num)
    
49.         x = self.dropout(x)
    
50.         logit = F.log_softmax(self.fc1(x), dim=1)
    
51.         return logit
    
52. 
    
53.     def init_weight(self):
    
54.         for m in self.modules():
    
55.             if isinstance(m, nn.Conv2d):
    
56.                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
    
57.                 m.weight.data.normal_(0, math.sqrt(2. / n))
    
58.                 if m.bias is not None:
    
59.                     m.bias.data.zero_()
    
60.             elif isinstance(m, nn.BatchNorm2d):
    
61.                 m.weight.data.fill_(1)
    
62.                 m.bias.data.zero_()
    
63.             elif isinstance(m, nn.Linear):
    
64.                 m.weight.data.normal_(0, 0.01)
    
65.                 m.bias.data.zero_()
    
66.                
    
67. 
    
68. textCNN_param = {
    
69. #     'vocab_size': len(word2ind),
    
70.     'vocab_size': 180,
    
71.     'embed_dim': 60,
    
72. #     'class_num': len(label_w2n),
    
73.     'class_num': 190,
    
74.     "kernel_num": 16,
    
75.     "kernel_size": [3, 4, 5],
    
76.     "dropout": 0.5,
    
77. }
    
78. dataLoader_param = {
    
79.     'batch_size': 128,
    
80.     'shuffle': True,
    
81. }
    
82. 
    
83. net = textCNN(textCNN_param)
    
84. print(net)

复制代码

输出结果: