Pytorch的学习——Yolov3的搭建

Yolov3的搭建

整个Yolov3一共经历三次下采样三次上采样，最后一共有三个输出，输出的维度分别为[通道数， 8， 8]， [通道数， 16， 16]， [通道数， 32， 32]

import torch.nn as nn


# yolo的darknet模块
class Darknet(nn.Module):
    def __init__(self, num_in, num_out, ksize, stride=1, padding=1):
        super(Darknet, self).__init__()
        self.conv1 = nn.Conv2d(num_in, num_out, ksize, stride, padding)
        self.bn1 = nn.BatchNorm2d(num_out)
        self.leakyReLU = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.leakyReLU(x)
        return x


# 残差模块 同resnet的残差块
class ResidualBlock(nn.Module):
    def __init__(self, num_in, num_out, num_block):
        super(ResidualBlock, self).__init__()
        self.num_block = num_block
        self.dark_conv1 = Darknet(num_in, num_out, ksize=3, stride=2, padding=1)
        self.dark_conv2 = []
        self.convblock = None
        for i in range(self.num_block):
            layers = list([])
            layers.append(Darknet(num_out, num_out//2, ksize=1, stride=1, padding=0))
            layers.append(Darknet(num_out//2, num_out, ksize=3, stride=1, padding=1))
            self.dark_conv2.append(nn.Sequential(*layers))
        self.dark_conv2 = nn.ModuleList(self.dark_conv2)

    def forward(self, x):
        x = self.dark_conv1(x)
        for convblock in self.dark_conv2:
            self.convblock = convblock
            residual = x
            x = self.convblock(x)
            x = x + residual
        return x


# darknet集成模块
class LastLayer(nn.Module):
    def __init__(self, num_in, num_out, num_out2):
        super(LastLayer, self).__init__()
        self.dark_conv1 = Darknet(num_in, num_out, ksize=1, stride=1, padding=0)
        self.dark_conv2 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.dark_conv3 = Darknet(num_out * 2, num_out, ksize=1, stride=1, padding=0)
        self.dark_conv4 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.dark_conv5 = Darknet(num_out * 2, num_out, ksize=1, stride=1, padding=0)

        self.dark_conv6 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.conv7 = nn.Conv2d(num_out * 2, num_out2, 1, stride=1, padding=0)

    def forward(self, x):
        x = self.dark_conv1(x)
        x = self.dark_conv2(x)
        x = self.dark_conv3(x)
        x = self.dark_conv4(x)
        x = self.dark_conv5(x)
        y = self.dark_conv6(x)
        y = self.conv7(y)
        return x, y

# yolov3主体
class Yolov3(nn.Module):
    def __init__(self, num_anchor, num_class):
        super(Yolov3, self).__init__()
        self.dark_conv1 = Darknet(3, 32, ksize=3, stride=1, padding=1)
        self.res1 = ResidualBlock(32, 64, 1)
        self.res2 = ResidualBlock(64, 128, 2)
        self.res3 = ResidualBlock(128, 256, 8)
        self.res4 = ResidualBlock(256, 512, 8)
        self.res5 = ResidualBlock(512, 1024, 4)

        self.last1 = LastLayer(1024, 512, num_anchor * (num_class + 5))

    def forward(self, x):
        x = self.dark_conv1(x)  
        x = self.res1(x) 
        x = self.res2(x) 
        x3 = self.res3(x) 
        x4 = self.res4(x3)  
        x5 = self.res5(x4)
        x, y1 = self.last1(x5)
        x = self.up1(x)
        x = torch.cat((x, x4), 1)
        x,y2 = self.last2(x)
        x = self.up2(x)
        x = torch.cat((x, x3), 1)
        x,y3 = self.last3(x)
        # y1的维度[通道数, 8, 8], y2的维度[通道数, 16, 16], y3的维度[通道数, 32, 32]
        return y1, y2, y3 

该模型没有全连接层需要自己加，根据输出的不同以及输出的维度自己添加

训练

在这边小编只需要维度为[通道数，8，8]的输出信息所以将上面的代码修改了一些

import torch
import torch.nn as nn


class Darknet(nn.Module):
    def __init__(self, num_in, num_out, ksize, stride=1, padding=1):
        super(Darknet, self).__init__()
        self.conv1 = nn.Conv2d(num_in, num_out, ksize, stride, padding)
        self.bn1 = nn.BatchNorm2d(num_out)
        self.leakyReLU = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.leakyReLU(x)
        return x


class ResidualBlock(nn.Module):
    def __init__(self, num_in, num_out, num_block):
        super(ResidualBlock, self).__init__()
        self.num_block = num_block
        self.dark_conv1 = Darknet(num_in, num_out, ksize=3, stride=2, padding=1)
        self.dark_conv2 = []
        self.convblock = None
        for i in range(self.num_block):
            layers = list([])
            layers.append(Darknet(num_out, num_out//2, ksize=1, stride=1, padding=0))
            layers.append(Darknet(num_out//2, num_out, ksize=3, stride=1, padding=1))
            self.dark_conv2.append(nn.Sequential(*layers))
        self.dark_conv2 = nn.ModuleList(self.dark_conv2)

    def forward(self, x):
        x = self.dark_conv1(x)
        for convblock in self.dark_conv2:
            self.convblock = convblock
            residual = x
            x = self.convblock(x)
            x = x + residual
        return x


class LastLayer(nn.Module):
    def __init__(self, num_in, num_out, num_out2):
        super(LastLayer, self).__init__()
        self.dark_conv1 = Darknet(num_in, num_out, ksize=1, stride=1, padding=0)
        self.dark_conv2 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.dark_conv3 = Darknet(num_out * 2, num_out, ksize=1, stride=1, padding=0)
        self.dark_conv4 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.dark_conv5 = Darknet(num_out * 2, num_out, ksize=1, stride=1, padding=0)

        self.dark_conv6 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.conv7 = nn.Conv2d(num_out * 2, num_out2, 1, stride=1, padding=0)
        self.fc = nn.Linear(8 * 8 * num_out2, 2)

    def forward(self, x):
        x = self.dark_conv1(x)
        x = self.dark_conv2(x)
        x = self.dark_conv3(x)
        x = self.dark_conv4(x)
        x = self.dark_conv5(x)
        y = self.dark_conv6(x)
        y = self.conv7(y)
        y = y.view(y.size(0), -1)
        y = self.fc(y)
        return x, y


class Yolov3(nn.Module):
    def __init__(self, num_anchor, num_class):
        super(Yolov3, self).__init__()
        self.dark_conv1 = Darknet(3, 32, ksize=3, stride=1, padding=1)
        self.res1 = ResidualBlock(32, 64, 1)
        self.res2 = ResidualBlock(64, 128, 2)
        self.res3 = ResidualBlock(128, 256, 8)
        self.res4 = ResidualBlock(256, 512, 8)
        self.res5 = ResidualBlock(512, 1024, 4)

        self.last1 = LastLayer(1024, 512, num_anchor * (num_class + 5))

    def forward(self, x):
        x = self.dark_conv1(x)  
        x = self.res1(x)  
        x = self.res2(x)  
        x3 = self.res3(x)
        x4 = self.res4(x3) 
        x5 = self.res5(x4) 
        x, y1 = self.last1(x5)  
        return y1

# 检验模型
if __name__ == '__main__':
    model = Yolov3(256, 2)
    input = torch.randn(1, 3, 256, 256)
    out = model(input)
    print(out.shape)

train模块

对模型进行训练与测试

import torch
import cv2
import torch.nn as nn
import torch.utils.data
from PIL import Image
import torchvision
import torchvision.transforms as transforms
from model.YoLov3 import Yolov3


# 超参数
EPOCH = 30
BATCH_SIZE = 10
LR = 0.001

IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')


def default_loader(path):
    with open(path, 'rb') as f:
        img = Image.fromarray(cv2.imread(f.name))
        return img


def load_train_data():
    path = "../medias/fire_pictures/train"  # 路径
    train_set = torchvision.datasets.ImageFolder(path,
                                                 transform=transforms.Compose([
                                                     transforms.Resize((256, 256)),
                                                     transforms.ToTensor()]),
                                                 loader=default_loader
                                                 )

    train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
    return train_loader


def load_test_data():
    path = "../medias/fire_pictures/test"
    test_set = torchvision.datasets.ImageFolder(path,
                                                transform=transforms.Compose([
                                                    transforms.Resize((256, 256)),
                                                    transforms.ToTensor()]),
                                                loader=default_loader
                                                )
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)
    return test_loader


def train(filename, model):
    train_loader = load_train_data()
    with torch.no_grad():
        model = model.cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
    loss_func = nn.CrossEntropyLoss()
    for epoch in range(EPOCH):
        for step, (b_x, b_y) in enumerate(train_loader):
            if torch.cuda.is_available():
                with torch.no_grad():
                    b_x = b_x.cuda()
                    b_y = b_y.cuda()
            output = model(b_x)
            loss = loss_func(output, b_y)
            optimizer.zero_grad()  
            loss.backward() 
            optimizer.step() 

            if step % 50 == 0:
                print('Epoch: ', epoch)
    cnn = model .cpu()
    torch.save(cnn, '../network/' + filename)


def test(filename):
    cnn = torch.load('./network/' + filename).cuda()
    test_loader = load_test_data()
    data_iter = iter(test_loader)
    images, labels = data_iter.__next__()
    outputs = cnn(images.cuda())
    predicted = torch.softmax(outputs, 1)
    p = torch.max(predicted, 1)
    print(p[1].cpu().numpy(), labels.cpu().numpy())
    j = 0
    for i in range(len(p[1].cpu().numpy())):
        a = int(p[1].cpu().numpy()[i]) - int(labels.cpu().numpy()[i])
        if a == 0:
            j += 1
    print(j / BATCH_SIZE)


if __name__ == '__main__':
    yolo = Yolov3(256, 2)
    train("yolov3_fire_256.pkl", yolo)