PYTORCH

import os
import itertools
from torchvision import transforms
from torch.utils.data import Dataset
import os
from PIL import Image
import torch
from torch.utils.data import DataLoader
import torch
import torch.nn as nn
from torchvision.datasets import mnist # 导入 pytorch 内置的 mnist 数据
from torch.autograd import Variable
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np
import torch
from torchvision.datasets import mnist # 导入 pytorch 内置的 mnist 数据
from torch import nn
from torch.autograd import Variable
import time
class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, (3, 3), (1, 1), 2), # out (32, 482, 446)
            nn.BatchNorm2d(32),
            nn.ReLU(inplace = True),
            nn.Conv2d(32, 64, (3, 3), (1, 1), 2),  # out (64, 482, 446)
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, (3, 3), (1, 1), 2),  # out (128, 482, 446)
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2) # out (128, 241, 223)
        )

        self.fc = nn.Sequential(
            nn.Linear(7058432, 300),
            nn.ReLU(True),
            nn.Linear(300, 256),
            nn.ReLU(True),
            nn.Linear(256, 3),
            nn.ReLU(True)
        )

    def forward(self, x):
        # in torch.Size([batchsize, 1, 28, 28])
        x = self.conv(x) # out torch.Size([batchsize, 128, 14, 14])
        x = x.view(x.size(0), -1) # out (batchsize, 128*14*14)
        x = self.fc(x) # out (batchsize, 10)
        return x

def concat(s, l):
    return([s + i for i in l])

mydic = {
'embedded_small':0, 'surface_small':1, 'un_visible':2
}
class MyData():
    def __init__(self, root_dir, label_dir = None, ):
        path_append = []
        self.img_path = []
        if label_dir is None:
            label_dir = ['embedded_small', 'surface_small', 'un_visible']
        self.root_dir = root_dir
        self.label_dir = label_dir
        for i in range(len(self.label_dir)):
            full_name = concat(self.root_dir + self.label_dir[i] + '/', os.listdir(os.path.join(self.root_dir + self.label_dir[i])))
            self.img_path.append(full_name)
        self.img_path = list(itertools.chain.from_iterable(self.img_path))

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = img_name
        img = Image.open(img_item_path) #PIL img
        img = self.data_tf(img)
        print(img.size)
        img = np.asarray(img) #array
        print(img.shape)
        label = mydic[str.split(img_name, '/')[2]]
        return img, label

    def __len__(self):
        return len(self.img_path)

    @staticmethod
    def data_tf(x):
        x = transforms.CenterCrop((482, 446))(x)  # PIL img
        x = np.array(x, dtype='float32') / 255
        x = (x - 0.5) / 0.5  # 标准化，这个技巧之后会讲到
        x = x.reshape((3, 482, 446))  # 拉平
        x = torch.from_numpy(x)
        return x

if __name__ == '__main__':
    dataset = MyData(root_dir='./dataset/')
    train_data = DataLoader(dataset, batch_size=5, shuffle=True)

    net = MyNet()
    # 定义 loss 函数
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(net.parameters(), 1e-1)  # 使用随机梯度下降，学习率 0.1

    # 开始训练
    losses = []
    acces = []
    eval_losses = []
    eval_acces = []

    for e in range(1):
        train_loss = 0
        train_acc = 0
        net.train()
        for im, label in train_data:
            im = Variable(im)

            label = Variable(label)
            print(label)
            # 前向传播

            out = net(im)
            print(out.shape)
            print(label.shape)
            loss = criterion(out, label)
            print(label.size)
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # 记录误差
            train_loss += loss.item()
            # 计算分类的准确率
            pred = out.max(1)
            num_correct = (pred == label).sum().item()
            acc = num_correct / im.shape[0]
            train_acc += acc

        losses.append(train_loss / len(train_data))
        acces.append(train_acc / len(train_data))
        # 在测试集上检验效果
        # eval_loss = 0
        # eval_acc = 0
        # net.eval()  # 将模型改为预测模式
        # for im, label in test_data:
        #     im = Variable(im)
        #     label = Variable(label)
        #     out = net(im)
        #     loss = criterion(out, label)
        #     print(loss.item)
        #     # 记录误差
        #     eval_loss += loss.item()
        #     # 记录准确率
        #     _, pred = out.max(1)
        #     num_correct = (pred == label).sum().item()
        #     acc = num_correct / im.shape[0]
        #     eval_acc += acc

        # eval_losses.append(eval_loss / len(test_data))
        # eval_acces.append(eval_acc / len(test_data))
        print('epoch: {}, Train Loss: {:.6f}, Train Acc: {:.6f}'
              .format(e, train_loss / len(train_data), train_acc / len(train_data)))