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=