通过Pytorch实现的AlexNet实现多分类项目留档

//训练代码(mytrain.py)：

    import os
    from torchvision import transforms
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image
    import numpy as np
    import torch
    from torch import optim
    import torch.nn as nn
    from torch.autograd import Variable
    from torchvision.models import alexnet

    import tqdm
    from tqdm import *


    def pre_process(path):
        # 传入训练集地址,预处理，地址表，打标签
        print('preprocessing the dataset in ', path)
        data = []
        real_list = os.listdir(path+'training_real/')
        fake_list = os.listdir(path+'training_fake/')
        # for jupyterlab
        if '.ipynb_checkpoints' in real_list : real_list.remove('.ipynb_checkpoints')
        if '.ipynb_checkpoints' in fake_list : fake_list.remove('.ipynb_checkpoints')
        # finished
        print('len of real_list:',len(real_list))
        print('len of fake_list:',len(fake_list))


        for imgpath in real_list:
            data.append([path+'training_real/'+imgpath,[0,0,0,0]])

        for imgpath in fake_list:
            lable_str = imgpath[-8:-4]
            if imgpath == '.ipynb_checkpoints' : print(imgpath,'in fake')
            lable = []
            for i in lable_str: 
                lable.append(int(i))
            data.append([path+'training_fake/'+imgpath,lable])

        print('len of data:',len(data))
        return data


    class MyDataset(Dataset):
        def __init__(self, data, transform, loader):
            self.data = data
            self.transform = transform
            self.loader = loader

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

        def __getitem__(self, idx):
            imgpath, label = self.data[idx]
            image = self.loader(imgpath)
            image = self.transform(image)
            return image, label


    def Myloader(path):
            return Image.open(path).convert('L').convert('RGB') # 转灰度图像


    def data_load(dataset_path):
        print('loading dataset......')
        # 返回 Dataloader_train

        # 自定义 transform
        transform = transforms.Compose([
            transforms.Resize((256,256)),
            transforms.ToTensor()
        ])

        # 预处理训练集集
        data = pre_process(dataset_path)

        # 初始化训练集
        train_data = MyDataset(data, transform=transform, loader=Myloader)
        Dtr = DataLoader(dataset=train_data, batch_size=50, shuffle=True, num_workers=0)

        return Dtr


    if __name__ == "__main__":
        epoch_num = 30
        best_model = None
        min_epochs = 5
        min_val_loss = 5
        epochs = 30
        lr = 0.0001
        batch_size = 25
        device = "cuda"

        dataset_path = './real_and_fake_face_training/'  # 训练集文件夹路径

        print(f'''
        paramater:
        ------------------
        epoch_num    \t {epoch_num} 
        best_model   \t {best_model}
        min_epochs   \t {min_epochs}
        min_val_loss   \t {min_val_loss}
        epochs   \t {epochs}
        lr      \t {lr}
        batch_size \t {batch_size}
        device   \t {device}
        dataset_path  \t {dataset_path}
        save_path  \t {save_path}
        ------------------
        ''')


        Dtr = data_load(dataset_path)
        model = alexnet(pretrained=True)
        model.classifier._modules['6'] = nn.Sequential(nn.Linear(4096, 4)) # 简单更改后的AlexNet
        print(model)

        print("training......")
        model = model.to(device)
        optimizer = optim.Adam(model.parameters(), lr=lr) # 优化方法 Adam
        criterion = nn.MultiLabelSoftMarginLoss().to(device)  # 损失函数用 MultiLabelSoftMarginLoss()

        for epoch in tqdm(range(epoch_num), ascii=True):
            train_loss = []
            for batch_idx, (data, target) in enumerate(Dtr, 0):

                target = torch.t(torch.stack(target))
                # print(target)
                # print(data)
                data, target = Variable(data).to(device), Variable(target).to(device)

                optimizer.zero_grad() # 梯度清0
                output = model(data)  # 前向传播 
                loss = criterion(output, target)  # 计算误差
                loss.backward() # 反向传播
                optimizer.step() # 更新参数
                train_loss.append(loss.cpu().item())

            tqdm.write('Epoch {:03d} train_loss {:.5f} '.format(epoch, np.mean(train_loss)))

        torch.save(model.state_dict(), 'model_weights.pth')

```python
//测试代码(mytest.py)

    import os
    from torchvision import transforms
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image
    import torch
    from torch.autograd import Variable
    from torchvision.models import alexnet
    import torch.nn as nn

    class MyDataset(Dataset):
        def __init__(self, data, transform, loader):
            self.data = data
            self.transform = transform
            self.loader = loader

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

        def __getitem__(self, idx):
            imgpath, label = self.data[idx]
            image = self.loader(imgpath)
            image = self.transform(image)
            return image, label


    def Myloader(path):
            return Image.open(path).convert('L').convert('RGB') # 转灰度图像


    def load_data(path):
        print('Loading data....')
        test_list = os.listdir(path)

        data = []
        for imgpath in test_list:
            data.append([path + imgpath,[0,0,0,0]])

        print('len of data:',len(data))

        transform = transforms.Compose([
            transforms.Resize((256,256)),
            transforms.ToTensor()
        ])

        test_dataset = MyDataset(data,transform,Myloader)
        Dts = DataLoader(dataset=test_dataset, batch_size=50, shuffle=False, num_workers=0)

        return Dts


    if __name__ == '__main__':

        device = 'cuda'
        test_path = './real_and_fake_face_testing/real_and_fake_face_testing/'

        model = alexnet()
        model.classifier._modules['6'] = nn.Sequential(nn.Linear(4096, 4))
        model.load_state_dict(torch.load('model_weights.pth'))
        model.eval()
        model.to(device)

        Dts = load_data(test_path)
        test_list = os.listdir(test_path)
        pred = []

        print('Testing...')
        for batch_idx, (data, lable) in enumerate(Dts, 0):
            lable = torch.t(torch.stack(lable))
            data, lable = Variable(data).to(device), Variable(lable).to(device)
            temp = model(data)
            # print(temp)
            pred.append(temp)

        # file=open('result.txt',mode='x')
        # print(pred)
        print('-----------------map-----------------')
        # for k in pred : file.write(str(k[0]))
        # file.close()

        res = []
        for k in pred:
            for z in k: res.append(z.cpu().detach().numpy().tolist())
        print('the res len:',len(res))
        for i in range(len(res)):
            print(test_list[i],res[i],sep='\t')