torch.nn.Module中常用函数

torch.nn.Module

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)