[Pytorch] PyTorch Dataloader数据读取以及训练实现过程

Resnet18实现猫狗分类！

本篇文章主要是本人入门Pytroch 的一些笔记。注意事项记录在代码之后！

import torch
torch.backends.cudnn.benchmark=True
import torch.nn as nn
import torch.utils.data as data
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
import torch.optim as optim
import torch.optim.lr_scheduler
import torch.nn.functional as F
import argparse
import time
import numpy as np
import sys
import cv2
import copy
import subprocess
from tqdm import tqdm
import os
from torch.optim import lr_scheduler

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

########## Hyperparams ######
num_epochs = 4
lr = 0.001
store_dir = '/home/hurw/Documents/program_all/cat_dog_pytorch/data/lianshou1.pth'

########## Loader #########
def load_dataset():
	#torchvision.transforms.Compose是用来管理所有transforms操作
	data_transforms = {
		'train': transforms.Compose([
			transforms.RandomResizedCrop(224),
			transforms.RandomHorizontalFlip(),
			transforms.ToTensor(),
			transforms.Normalize(mean,std)
		]),
		'val': transforms.Compose([
			transforms.Scale(256),
			transforms.CenterCrop(224),
			transforms.ToTensor(),
			transforms.Normalize(mean,std)
		]),
		'test': transforms.Compose([
			transforms.Scale(256),
			transforms.CenterCrop(224),
			transforms.ToTensor(),
			transforms.Normalize(mean,std)
		])
	}

	data_dir = '/home/hurw/Documents/program_all/cat_dog_pytorch/data'
	image_datasets = {x : datasets.ImageFolder(os.path.join(data_dir, x),
					  data_transforms[x])for x in ['train', 'val', 'test']}



	data_loaders = {x: data.DataLoader(image_datasets[x],
					batch_size=32,num_workers=12, shuffle=True)
					for x in ['train', 'val', 'test']}

	# data_loaders = {x : torch.utils.data.DataLoader(image_datasets[x], batch_size=32,
	# 											num_workers=12, shuffle=True)
	# 					for x in ['train', 'val', 'test']}
	data_size = {x : len(image_datasets[x]) for x in ['train', 'val', 'test']}
	return data_loaders, data_size

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = models.resnet18(pretrained=False)
        for param in self.model.parameters():
            param.requires_grad = False # 我们只是需要训练好的参数，不需要梯度信息，故设置为False
            #param.requires_grad = True

        self.model.fc = nn.Linear(self.model.fc.in_features,2,bias=False)
    def forward(self, x):
        x = self.model(x)
        return x


def train(data_loader, data_size):
	model = Model()
	model = model.cuda()  # Moves all model parameters and buffers to the GPU.
	criterion = nn.CrossEntropyLoss()
	# 这种写法只能训练最后的全连接层
	#optimizer = optim.SGD(model.model.fc.parameters(), lr=lr, momentum=