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

最新推荐文章于 2025-06-24 16:02:16 发布

为自己勇敢

最新推荐文章于 2025-06-24 16:02:16 发布

阅读量2.5w

点赞数 36

CC 4.0 BY-SA版权

分类专栏：机器学习

本文链接：https://blog.youkuaiyun.com/Summer_And_Opencv/article/details/88973713

本文介绍了使用PyTorch的Dataloader和Resnet18实现猫狗分类的过程，详细讲解了数据集结构、数据读取、数据预处理、模型构建、损失函数和优化器的设定，以及训练流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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.param