初始化

首先，介绍我们导入的包和基础的网络结构

import torch
import torch.nn as nn

#可替代网络结构部分
'''
神经网络类的定义
    1. 输入卷积: in_channel = 1, out_channel = 12, kernel_size = (5, 5), stride = (2, 2), padding = 2
    2. 激活函数: 1.7159Tanh(2/3*x)
    3. 第二层卷积: in_channel = 12, out_channel = 12, kernel_size = (5, 5), stride = (2, 2), padding = 2
    4. 激活函数同上
    5. 全连接层: 192 * 30
    6. 激活函数同上
    7. 全连接层：30 * 10
    8. 激活函数同上

    按照论文的说明，需要对网络的权重进行一个[-2.4/F_in, 2.4/F_in]的均匀分布的初始化
'''


class LeNet1989(nn.Module):
    def __init__(self):
        super(LeNet1989, self).__init__()

        self.conv1 = nn.Conv2d(1, 12, 5, stride=2, padding=2)
        self.act1 = nn.Tanh()
        self.conv2 = nn.Conv2d(12, 12, 5, stride=2, padding=2)
        self.act2 = nn.Tanh()

        self.fc1 = nn.Linear(192, 30)
        self.act3 = nn.Tanh()
        self.fc2 = nn.Linear(30, 10)
        self.act4 = nn.Tanh()

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                F_in = m.kernel_size[0] * m.kernel_size[1] * m.in_channels
                m.weight.data = torch.rand(m.weight.data.size()) * 4.8 / F_in - 2.4 / F_in
            if isinstance(m, nn.Linear):
                F_in = m.in_features
                m.weight.data = torch.rand(m.weight.data.size()) * 4.8 / F_in - 2.4 / F_in

    def forward(self, x):
        x = self.conv1(x)
        x = 1.7159 * self.act1(2.0 * x / 3.0)

        x = self.conv2(x)
        x = 1.7159 * self.act2(2.0 * x / 3.0)

        x = x.view(-1, 192)

        x = self.fc1(x)
        x = 1.7159 * self.act3(2.0 * x / 3.0)

        x = self.fc2(x)
        out = 1.7159 * self.act4(2.0 * x / 3.0)

        return out

查看网络参数

方法一

推荐使用方法一，因为方法二可以得到更多的信息，但是要注意的是这一段看情况添加

看网络是在GPU跑还是CPU跑，我相信大部分是GPU，就用下面这个就可

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = LeNet1989().to(device)

用CPU就用这个

net = LeNet1989()

下面的部分就注意输入图片的input_size就可

from torchsummary import summary
#batch_size可以不指定，默认为-1，输入 模型（model）、输入尺寸（input_size）、批次大小（batch_size）、运行平台（device）
#summary(model, input_size, batch_size, device)
summary(net, input_size=(1, 16, 16), batch_size=1)

运行结果：

方法二

代码：

#查看网络参数部分
net = LeNet1989()
print(net)
params = list(net.parameters())
k = 0
for i in params:
    l = 1
    # print("该层的结构："+ str(list(i.size())))
    for j in i.size():
        l *= j
    # print("该层参数和："+ str(l))
    k = k + l
print('总参数和：' + str(k))

运行结果：

查看网络结构

from tensorboardX import SummaryWriter  # pip install tensorboardX
x = torch.randn(1, 16, 16) # 例如：x = torch.randn(1, 64, 64, 64)
net = LeNet1989()
writer = SummaryWriter('./tensorboard')
with writer:
    writer.add_graph(net, (x,))
# 终端中输入：tensorboard --logdir=tensorboard/ --host=127.0.0.1

图片通道转换

彩色图转灰度图

from torchvision import transforms

transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1), # 彩色图像转灰度图像num_output_channels默认1
    transforms.ToTensor()
])

彩色图（三通道）转指定R,G,B通道

def change_image_channels(image):
    # 3通道转单通道
    if image.mode == 'RGB':
        r, g, b = image.split()
    return r,g,b

模型测试调用

.pt, .pth, .pkl文件都可

保存：

torch.save(model, mymodel.pth)#保存整个model的状态

调用：

model=torch.load(mymodel.pth)#这里已经不需要重构模型结构了，直接load就可以
model.eval()