pyTorch学习2:搭建网络

利用pytorch搭建网络的核心在于继承类torch.nn.Module并重写__init__()和forward()函数，添加新层的方式有如下几种：

1. 直接添加

self.conv1 = torch.nn.Conv2d(...)
self.fc1 = torch.nn.Linear(...)

2. add_module方法

self.add_module("conv1", torch.nn.Conv2d(...))
self.add_module("fc1", torch.nn.Linear(...))

3. Sequential方法

self.conv1 = Sequential(
		torch.nn.Conv2d(...),
		torch.nn.ReLU(),
		torch.nn.MaxPool2d(...)
		)

4. ModuleList方法
   把要添加的层写入一个list，随后使用ModuleList方法添加至网络

self.linears = torch.nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

需要注意的是，使用这种方法搭建的网络，在forward时需逐层前传：

for layer in self.linears:
	x = layer(x)

CNN

cnn.py

import torch
import torch.nn as nn
import torch.nn.functional as F


class CustomCNN(nn.Module):
    def __init__(self):
        super(CustomCNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=7*7*32, out_features=256),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(in_features=256, out_features=10),
        )
    def forward(self, x):
        x = self.conv1(x)
        # print("conv1 {}".format(x.shape))
        x = self.conv2(x)
        # print("conv2 {}".format(x.shape))
        x = x.view(x.size(0), -1)
        # print("view {}".format(x.shape))
        x = self.fc1(x)
        # print("fc1 {}".format(x.shape))
        x = self.fc2(x)
        # print("fc2 {}".format(x.shape))
        return x

train_cnn.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from mnist.cnn import CustomCNN

train_set = MNIST(root="../data", train=True, transform=ToTensor())
test_set = MNIST(root="../data", train=True, transform=ToTensor())

train_loader = DataLoader(train_set, batch_size=16, shuffle=True, drop_last=False)
test_loader = DataLoader(test_set, batch_size=16, shuffle=True, drop_last=False)

net = CustomCNN()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=net.parameters(), lr=0.001)

for epoch in range(10):
    o_loss = 0
    o_step = 0
    for step, (data, label) in enumerate(train_loader):
        out = net(data)
        loss = criterion(out, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        o_loss += loss
        o_step += 1
    print("epoch {} loss {}".format(epoch + 1, o_loss/o_step))
    tp = 0
    v_size = 0
    v_loss = 0
    for _, (v_data, v_label) in enumerate(test_loader):
        v_out = nn.functional.softmax(net(v_data), dim=1)
        v_loss += criterion(v_out, v_label)
        v_pred = torch.max(v_out, dim=1)[1]
        tp += sum(v_pred == v_label)
        v_size += v_pred.shape[0]

    print("val acc={}/{}={}".format(tp.numpy(), v_size, tp.numpy()/v_size))
    print("val loss={}".format(v_loss))

多层感知机

perceptron.py

import torch
import torch.nn as nn
import torch.nn.functional as F


class CustomPerceptron(nn.Module):
    def __init__(self, config):
        super(CustomPerceptron, self).__init__()
        layer_list = []
        for l in range(len(config)):
            if l == 0:
                layer_list.append(nn.Sequential(
                    nn.Linear(in_features=784, out_features=config[l]),
                    nn.ReLU()
                ))
            else:
                layer_list.append(nn.Sequential(
                    nn.Linear(in_features=config[l - 1], out_features=config[l]),
                    nn.ReLU()
                ))
        layer_list.append(nn.Linear(in_features=config[-1], out_features=10))
        self.network = nn.ModuleList(layer_list)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        for layer in self.network:
            x = layer(x)
        return x

train_perceptron.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from mnist.perceptron import CustomPerceptron


class Trainer:
    def __init__(self, net, train_loader, test_loader):
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.net = net
        print(self.net)

    def train(self, criterion=nn.CrossEntropyLoss, optimizer=torch.optim.Adam, lr=0.001, max_epoch=10, eval_interval=1):
        self.criterion = criterion()
        self.optimizer = optimizer(self.net.parameters(), lr=lr)
        # self.criterion = nn.CrossEntropyLoss()
        # self.optimizer = torch.optim.Adam(self.net.parameters(), lr=0.001)
        for epoch in range(max_epoch):
            oloss = 0
            ostep = 0
            for step, (data, label) in enumerate(train_loader):
                out = self.net(data)
                loss = self.criterion(input=out, target=label)
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                oloss += loss
                ostep += 1
            print("epoch {} train loss {}".format(epoch + 1, oloss / ostep))
            if (epoch + 1) % eval_interval == 0:
                l, acc = self.test(self.test_loader)
                print("epoch {} val loss {} val acc {}".format(epoch + 1, l, acc))

    def test(self, data_loader):
        oloss = 0
        nsample = 0
        tp = 0
        for step, (data, label) in enumerate(data_loader):
            out = self.net(data)
            oloss += self.criterion(input=out, target=label)
            out = nn.functional.softmax(out, dim=1)
            pred = torch.max(out, dim=1)[1]
            tp += sum(pred == label)
            nsample += pred.shape[0]
        return oloss, tp.numpy()/nsample


if __name__ == "__main__":
    train_set = MNIST(root="../data", train=True, transform=ToTensor())
    test_set = MNIST(root="../data", train=False, transform=ToTensor())

    train_loader = DataLoader(train_set, batch_size=16, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_set, batch_size=16, shuffle=True, drop_last=True)

    config = [256, 64]
    net = CustomPerceptron(config)

    trainer = Trainer(net=net, train_loader=train_loader, test_loader=test_loader)
    trainer.train()

参考文献：添加链接描述