0023-pytorch入门-读取自己的数据集并进行分类脚本

刺骨的言语ヽ痛彻心扉 2023-01-08 06:23 218阅读 0赞

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/19 12:06
    # @Author  : Johnson
    
    from __future__ import print_function, division
    
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.optim import lr_scheduler
    from torch.autograd import Variable
    from torchvision import  models, transforms
    
    import time
    import os
    from torch.utils.data import Dataset
    
    from PIL import Image
    
    # use PIL Image to read image
    def default_loader(path):
        try:
            img = Image.open(path)
            return img.convert('RGB')
        except:
            print("Cannot read image: {}".format(path))
    
    # define your Dataset. Assume each line in your .txt file is [name/tab/label], for example:0001.jpg 1
    class customData(Dataset):
        def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader):
            with open(txt_path) as input_file:
                lines = input_file.readlines()
                self.img_name = [os.path.join(img_path, line.strip().split('\t')[0]) for line in lines]
                self.img_label = [int(line.strip().split('\t')[-1]) for line in lines]
            self.data_transforms = data_transforms
            self.dataset = dataset
            self.loader = loader
    
        def __len__(self):
            return len(self.img_name)
    
        def __getitem__(self, item):
            img_name = self.img_name[item]
            label = self.img_label[item]
            img = self.loader(img_name)
    
            if self.data_transforms is not None:
                try:
                    img = self.data_transforms[self.dataset](img)
                except:
                    print("Cannot transform image: {}".format(img_name))
            return img, label
    
    def train_model(model, criterion, optimizer, scheduler, num_epochs, use_gpu):
        since = time.time()
    
        best_model_wts = model.state_dict()
        best_acc = 0.0
    
        for epoch in range(num_epochs):
            begin_time = time.time()
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            print('-' * 10)
    
            # Each epoch has a training and validation phase
            for phase in ['train', 'val']:
    	    count_batch = 0
                if phase == 'train':
                    scheduler.step()
                    model.train(True)  # Set model to training mode
                else:
                    model.train(False)  # Set model to evaluate mode
    
                running_loss = 0.0
                running_corrects = 0.0
    
                # Iterate over data.
                for data in dataloders[phase]:
                    count_batch += 1
                    # get the inputs
                    inputs, labels = data
    
                    # wrap them in Variable
                    if use_gpu:
                        inputs = Variable(inputs.cuda())
                        labels = Variable(labels.cuda())
                    else:
                        inputs, labels = Variable(inputs), Variable(labels)
    
                    # zero the parameter gradients
                    optimizer.zero_grad()
    
                    # forward
                    outputs = model(inputs)
                    _, preds = torch.max(outputs.data, 1)
                    loss = criterion(outputs, labels)
    
                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                    # statistics
                    running_loss += loss.data[0]
                    running_corrects += torch.sum(preds == labels.data).to(torch.float32)
    
                    # print result every 10 batch
                    if count_batch%10 == 0:
                        batch_loss = running_loss / (batch_size*count_batch)
                        batch_acc = running_corrects / (batch_size*count_batch)
                        print('{} Epoch [{}] Batch [{}] Loss: {:.4f} Acc: {:.4f} Time: {:.4f}s'. \
                              format(phase, epoch, count_batch, batch_loss, batch_acc, time.time()-begin_time))
                        begin_time = time.time()
    
                epoch_loss = running_loss / dataset_sizes[phase]
                epoch_acc = running_corrects / dataset_sizes[phase]
                print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
    
                # save model
                if phase == 'train':
                    if not os.path.exists('output'):
                        os.makedirs('output')
                    torch.save(model, 'output/resnet_epoch{}.pkl'.format(epoch))
    
                # deep copy the model
                if phase == 'val' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    best_model_wts = model.state_dict()
    
        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(
            time_elapsed // 60, time_elapsed % 60))
        print('Best val Acc: {:4f}'.format(best_acc))
    
        # load best model weights
        model.load_state_dict(best_model_wts)
        return model
    
    if __name__ == '__main__':
    
        data_transforms = {
            'train': transforms.Compose([
                transforms.RandomSizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ]),
            'val': transforms.Compose([
                transforms.Scale(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ]),
        }
    
        use_gpu = torch.cuda.is_available()
    
        batch_size = 32
        num_class = 2
    
        image_datasets = {x: customData(img_path='/ImagePath',
                                        txt_path=('/TxtFile/' + x + '.txt'),
                                        data_transforms=data_transforms,
                                        dataset=x) for x in ['train', 'val']}
    
        # wrap your data and label into Tensor
        dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
                                                     batch_size=batch_size,
                                                     shuffle=True) for x in ['train', 'val']}
    
        dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
    
        # get model and replace the original fc layer with your fc layer
        model_ft = models.resnet50(pretrained=True)
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Linear(num_ftrs, num_class)
    
        # if use gpu
        if use_gpu:
            model_ft = model_ft.cuda()
    
        # define cost function
        criterion = nn.CrossEntropyLoss()
    
        # Observe that all parameters are being optimized
        optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.005, momentum=0.9)
    
        # Decay LR by a factor of 0.2 every 5 epochs
        exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.2)
    
        # multi-GPU
        model_ft = torch.nn.DataParallel(model_ft, device_ids=[0,1])
    
        # train model
        model_ft = train_model(model=model_ft,
                               criterion=criterion,
                               optimizer=optimizer_ft,
                               scheduler=exp_lr_scheduler,
                               num_epochs=25,
                               use_gpu=use_gpu)
    
        # save best model
        torch.save(model_ft,"output/best_resnet.pkl")