Pytorch neural network construction template

Let's use this code to explain the process of neural network construction

import torch
import torchvision
from torch import nn
from torch.nn import MaxPool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
dataset = torchvision.datasets.CIFAR10("CIFAR10",train=False,transform=torchvision.transforms.ToTensor(),download=True)
#  Be careful dataset in transform The parameter receives an object , So we need to add parentheses , In addition, when using neural network for operation, the data type required is tensor type , therefore transforms Parameters to add .
dataloader = DataLoader(dataset,batch_size=64)

#  Building neural networks 
class Booze(nn.Module):

    #  Inherit nn.Module The initialization 
    def __init__(self):
        super(Booze, self).__init__()
        self.maxpool1 = MaxPool2d(kernel_size=3,ceil_mode=True)

    #  rewrite forward function 
    def forward(self,x):
        output = self.maxpool1(x)
        return output


#  Generate a neural network object 
obj = Booze()

writer = SummaryWriter("logs")
step = 0
for data in dataloader:
    imgs,targets = data
    writer.add_images("input",imgs,step)
    output = obj(imgs)
    writer.add_images("output",output,step)
    step+=1

writer.close()

The process can be divided into three parts ：

• 1. Data reading
• 2. Building neural networks
• 3. Use tensorboard visualization

1. Data reading

dataset = torchvision.datasets.CIFAR10("CIFAR10",train=False,transform=torchvision.transforms.ToTensor(),download=True)
#  Be careful dataset in transform The parameter receives an object , So we need to add parentheses , In addition, when using neural network for operation, the data type required is tensor type , therefore transforms Parameters to add .
dataloader = DataLoader(dataset,batch_size=64)

2. Building neural networks

#  Building neural networks 
class Booze(nn.Module):

    #  Inherit nn.Module The initialization 
    def __init__(self):
        super(Booze, self).__init__()
        self.maxpool1 = MaxPool2d(kernel_size=3,ceil_mode=True)

    #  rewrite forward function 
    def forward(self,x):
        output = self.maxpool1(x)
        return output

Inherit nn.Module, Inherit nn.Module The initialization of plus their own initialization , rewrite forword Method .

3. Use tensorboard visualization

#  Generate a neural network object 
obj = Booze()

writer = SummaryWriter("logs")
step = 0
for data in dataloader:
    imgs,targets = data
    writer.add_images("input",imgs,step)
    output = obj(imgs)
    writer.add_images("output",output,step)
    step+=1

writer.close()