可以通过torch.nn.Sequential来构建模型,将各层按顺序传入

一、代码

importtorchimporttorch.nnasnnfromtorchsummaryimportsummary# 1.定义数据x=torch.randn(10,3)# 2. 构建模型model=nn.Sequential(nn.Linear(in_features=3,out_features=4),nn.Tanh(),nn.Linear(in_features=4,out_features=4),nn.ReLU(),nn.Linear(in_features=4,out_features=2),nn.Softmax(dim=1),)# 定义一个参数初始化的函数definit_params(layer):ifisinstance(layer,nn.Linear):nn.init.xavier_uniform_(layer.weight)nn.init.constant_(layer.bias,val=0.1)# 3. 参数初始化model.apply(init_params)# 4. 前向传播output=model(x)print(output)summary(model,input_size=(3,),batch_size=10,device='cpu')#打印如下tensor([[0.2617,0.7383],[0.4800,0.5200],[0.2728,0.7272],[0.3490,0.6510],[0.4847,0.5153],[0.5021,0.4979],[0.1817,0.8183],[0.4712,0.5288],[0.2881,0.7119],[0.4895,0.5105]],grad_fn=<SoftmaxBackward0>)----------------------------------------------------------------Layer(type)Output Shape Param#================================================================Linear-1[10,4]16Tanh-2[10,4]0Linear-3[10,4]20ReLU-4[10,4]0Linear-5[10,2]10Softmax-6[10,2]0================================================================Total params:46Trainable params:46Non-trainable params:0----------------------------------------------------------------Input size(MB):0.00Forward/backwardpasssize(MB):0.00Params size(MB):0.00Estimated Total Size(MB):0.00----------------------------------------------------------------