import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
# 定义网络时需要继承nn.Module并实现它的forward方法,将网络中具有可学习参数的层放在构造函数__init__中
# 不具有可学习参数的层(如ReLu)既可以放在构造函数中也可以不放
class helloNet(nn.Module):
def __init__(self):
# nn.Module子类的函数必须在构造函数中执行父类的构造函数
super(helloNet, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
# 只要在nn.Module的子类中定义了forward函数,backward函数就会被使用Autograd自动实现
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=(2,2))
x = F.max_pool2d(F.relu(self.conv2(x)), kernel_size=(2,2))
x = x.view(x.size()[0], -1) # reshape, -1标识自适应
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# 创建网络
net = helloNet()
# 打印网络结构
print(net)
print(net.conv1.weight) # 打印某一层的参数
print(net.state_dict().keys()) # 打印所有参数名称
print("网络是训练模式还是推理模式: ", net.training) # training (bool) – Boolean represents whether this module is in training or evaluation mode.
print("可学习参数个数:")
print(len(list(net.parameters())))
# 打印所有参数及其尺寸
for param in net.parameters():
print(param, '->', param.size())
# 打印每个网络参数名称及尺寸
for name, parameters in net.named_parameters():
print(name, ":", parameters.size())
# 设置输入,并进行一次前向推理
input = Variable(torch.randn(1,1,32,32))
output = net(input) # 会调用自定义的forward函数
print(output.size())
net.zero_grad() # 将所有参数的梯度清零
output.backward(Variable(torch.ones(1,10))) # 反向传播
# 打印每一层的网络结构,modules()函数的作用是:Return an iterator over all modules in the network.
for idx, m in enumerate(net.modules()):
print(idx, '->', m)
# named_children()作用:Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
for name, module in net.named_children():
if name in ['conv2', 'fc2']:
print(module)
# named_modules()函数返回网络每层的名称及该层网络
for idx, m in enumerate(net.named_modules()):
print(idx, '->', m)
# 将网络放到gpu或gpu
gpu0 = torch.device("cuda:0")
cpu = torch.device("cpu")
net.to(gpu0, dtype=torch.half, non_blocking=True)
net.to(cpu)
torch.nn.Module中常用函数
最新推荐文章于 2024-09-02 23:14:21 发布