pytorch训练网络

import os
import numpy as np
import torch
import torch.nn as nn
import torch.functional as F
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
from torch.utils.data import DataLoader
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.autograd import Variable


captcha_array = list("0123456789abcdefghijklmnopqrstuvwxyz")

def text2vec(x):
    zeros = torch.zeros((4, 36), dtype = torch.long)
    for i in range(len(x)):
        zeros[i, captcha_array.index(x[i])] = 1
    vec = zeros
    return vec

def vec2text(x):
    label = ''
    label_tensor = torch.argmax(x, dim=1)
    for i in label_tensor:
        label += captcha_array[i.data]
    return label

def single_vec_2_text(label_tensor):
    label = ''
    for i in label_tensor:
        label += captcha_array[i.data]
    return label

def batch_vec_2_text(x):
    label_batch = []
    label_all = x.argmax(dim=2)
    for index_image in range(label_all.size(0)):
        label_batch.append(single_vec_2_text(label_all[index_image]))
    return label_batch

def compare_list(x, y):
    result_list = []
    for le in range(len(x)):
        if x[le] == y[le]:
            result_list.append(True)
    return result_list.count(True), len(x), result_list.count(True)/len(x)

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, (3, 3), (1, 1), 1), # in torch.Size([batch, 1, 60, 160])
            # nn.BatchNorm2d(32),
            nn.ReLU(inplace = True),
            nn.MaxPool2d(2) # out (batch, 32, 30, 80)
        )

        self.layer2 = nn.Sequential(
            nn.Conv2d(32, 64, (3, 3), (1, 1), 1),  # in torch.Size([
            # nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)  # out (batch, 64, 15, 40)
        )

        self.layer3 = nn.Sequential(
            nn.Conv2d(64, 128, (3, 3), (1, 1), 1),  # in torch.Size([
            # nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)  # out (batch, 128, 7, 20)
        )

        self.layer4 = nn.Sequential(
            nn.Conv2d(128, 256, (3, 3), (1, 1), 1),  # in torch.Size(
            # nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)  # out (batch, 256, 3, 10)
        )

        self.layer5 = nn.Sequential(
            nn.Linear(256 * 3 * 10, 1024),
            nn.Dropout(),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 4 * 36),
        )

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = out.view(out.size(0), -1) # 保留第一维度

        out = self.layer5(out)
        return out

class MyData(Dataset):
    def __init__(self, if_train = True):
        super(MyData, self).__init__()
        self.transforms = transforms.Compose(
        [
            transforms.Resize((60, 160)),
            transforms.Grayscale(),
            transforms.ToTensor()
        ]
    )
        if if_train == True:
            self.root = "./dataset5/train/"
        else:
            self.root = "./dataset5/test/"

    def __len__(self):
        return(len(os.listdir(self.root)))

    def __getitem__(self, item):
        img_path = self.root + os.listdir(self.root)[item]
        img = Image.open(img_path)
        img = self.transforms(img)
        label = os.listdir(self.root)[item].split('_')[0]
        label = text2vec(label)
        label = label.view(1, -1)[0] # 极其重要
        # t = torch.zeros(100, 1, 4, 36)
        # print(t.view(1, -1))
        # print(t.view(1, -1)[0])
        return img, label # img torch.Size([1, 60, 160])

if __name__ == '__main__':
    # 开始训练

    m = MyNet()
    # 定义 loss 函数
    loss_fn = nn.MultiLabelSoftMarginLoss()
    optimizer = torch.optim.Adam(m.parameters(), lr=0.0008)

    total_step = 0
    count = 80000
    count_a = 0

    for epoch in range(80000):
        train_data = MyData()
        test_data = MyData(if_train=False)
        # 使用 pytorch 自带的 DataLoader 定义一个数据迭代器
        train_dataloader = DataLoader(train_data, batch_size=10, shuffle=True)
        test_dataloader = DataLoader(test_data, batch_size=10, shuffle=False)

        for imgs, targets in train_dataloader:
            imgs = Variable(imgs)
            count_a += 1
            if count_a < count:
                m.train()
                outputs = m(imgs)

                loss = loss_fn(outputs, targets)
                outputs = outputs.view(-1, 4, 36)
                # print('预测为', batch_vec_2_text(outputs))
                targets = targets.view(-1, 4, 36)
                # print('实际为', batch_vec_2_text(targets))
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                total_step += 1
                print("训练{}次,loss:{}".format(total_step * 1, loss.item()),compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets)))
                with open('log.txt', 'a+') as f:
                    f.write("训练{}次,loss:{}".format(total_step * 1, loss.item()) + str(compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets))) + "\n")
                    f.close()
            else:
                break

            imgs, targets = next(iter(test_dataloader))
            imgs = Variable(imgs)
            m.eval()
            outputs = m(imgs)
            loss = loss_fn(outputs, targets)
            outputs = outputs.view(-1, 4, 36)
            # print('样本预测为', batch_vec_2_text(outputs))
            targets = targets.view(-1, 4, 36)
            # print('样本实际为', batch_vec_2_text(targets))
            print("训练{}次,样本loss:{}".format(total_step * 1, loss.item()),
                  compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets)))

            with open('log1.txt', 'a+') as f:
                f.write("训练{}次,loss:{}".format(total_step * 1, loss.item()) + str(
                    compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets))) + "\n")
                f.close()
            if count_a%15000 == 0:
                torch.save(m, "model.pth")

    torch.save(m, "model.pth")