pytorch学习之路——06

迁移学习、数据集划分

import torch
from torch import nn,optim
from torchvision import transforms,models
from torch.utils.data.dataset import Dataset
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler

import numpy as np
import pandas as pd
from PIL import Image
from collections import OrderedDict

# 加载csv文件
class DatasetFromCSV(Dataset):
    def __init__(self,path,h,w,train=None,transforms=None):
        self.data = pd.read_csv(path)
        self.labels = np.array(self.data.iloc[:,0])
        self.height = h
        self.width = w
        self.train = train
        self.transforms = transforms

    def __getitem__(self,idx):
        if self.train is not False:
            single_image_label = self.labels[idx]
            img_as_np = np.asarray(self.data.iloc[idx][1:]).reshape(self.height,self.width).astype(float)
            img_as_img = Image.fromarray(img_as_np)
            img_as_img = img_as_img.convert('RGB')
            if self.transforms is not None:
                img_as_tensor = self.transforms(img_as_img)
            return (img_as_tensor,single_image_label)
        else:
            imgs_as_np = np.asarray(self.data.iloc[idx][0:]).reshape(self.height,self.width).astype(f