pytorch-网络模型定义

网络模型定义

构建网络模型类
1.该类必须继承torch.nn.Module，并且在初始化方法中引入父类初始化方法
2.在__init__()中定义并初始化网络
3.在forward()中定义前向传递(反向传播根据前向传递自动定义)
使用时直接初始化该类，传入输入即可的得到输出

"""模块导入"""
import torch
import torch.nn as nn # 包含模型设计的模块和类
import torch.nn.functional as F 

"""网络模型定义"""
class Net(nn.Module):
	# 定义并初始化网络
	def __init__(self):
		# 继承父类初始化方法
		super(Net,self).__init__()
		# (1,32,3,1)->(in_channel,out_channel,kernel_size,batch)
		self.conv1 = nn.Conv2d(1,32,3,1)
		self.conv2 = nn.Conv2d(32,64,3,1)
		# dropout_rate=0.25
		self.dropout1 = nn.Dropout2d(0.25)
		self.dropout2 = nn.Dropout2d(0.5)
		# (9216,128)->(in_feature,out_feature)
		self.fc1 = nn.Linear(9216,128)
		self.fc2 = nn.Linear(128,10)
	
	# 指定前向传递
	def forward(self,x):
		x = self.conv1(x)
		x = F.relu(x)
		x = self.conv2(x)
		x = F.relu(x)
		x = F.max_pool2d(x,2)
		x = self.dropout1(x)
		x = torch.flatten(x,1)
		x = self.fc1(x)
		x = F.relu(x)
		x = self.dropout2(x)
		x = self.fc2(x)
		output = F.log_softmax(x,dim=1)
		return output

# 使用
input = torch.rand((1,1,28,28))
model = Net()
result = model(input)