pytorch之nn.Conv2d用法

nn.functional.Conv2d提供了卷积的计算函数，需要自己定义卷积核，输入图像的相关信息（值），从而帮助计算卷积。而nn.Conv2d类仅需要提供输入层，卷积核的尺寸信息即可计算，自动初始化权重等信息。

以下为nn.Conv2d使用的简单示例。

# 这部分介绍下卷积层的实现，以2维卷积层为例，位于nn.Conv2d
# 注意nn.functional.Conv2d是一个函数，而这里是卷积层是一个类
# 以下代码定义了一个运用卷积层Conv2d的类的简单示例，只涉及预测（正向传播过程），而没有学习过程（反向传播）。
import torch
import torchvision.datasets
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


class My_nn(nn.Module):
    def __init__(self):
        super(My_nn, self).__init__()
        # 定义卷积层，用于后面的forword函数。该层自动初始化权重向量的值。
        self.conv2d = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, padding=0, stride=1, bias=True)

    def forward(self, x):
        out = self.conv2d(x)
        return out


if __name__ == '__main__':
    # 定义数据集
    test_data = torchvision.datasets.CIFAR10('./dataset', train=False, transform=torchvision.transforms.ToTensor(),
                                             download=True)
    Data_Ld = DataLoader(test_data, batch_size=64, drop_last=True)  # DataLoader将数据打包，即组成mini-batch
    # 实例化类对象
    my_nn = My_nn()

    # 对测试数据集进行操作---只预测结果
    idx = 0
    writer = SummaryWriter("nn_conv2d")
    for data in Data_Ld:
        imgs, targets = data
        writer.add_images('original', imgs, idx)
        out_imgs = my_nn(imgs)  # 相当于out_imgs = my_nn.forward(imgs) 即正向传播。
        # print(imgs.shape)  # (64, 3, 32, 32)
        # print(out_imgs.shape)  # (64, 6, 30, 30)
        # 由于writer.add_images只能显示3通道的图像，因此要用reshape对结果图out_imgs像进行转换
        out_imgs = torch.reshape(out_imgs, (-1, 3, 30, 30))  # 将通道数改成3，而min-batch数可通过填-1让系统自动计算。
        writer.add_images('processed', out_imgs, idx)
        idx += 1

    writer.close()