pytorch一步一步在VGG16上训练自己的数据集

准备数据集及加载,ImageFolder

在很多机器学习或者深度学习的任务中,往往我们要提供自己的图片。也就是说我们的数据集不是预先处理好的,像mnist,cifar10等它已经给你处理好了,更多的是原始的图片。比如我们以猫狗分类为例。在data文件下,有两个分别为train和val的文件夹。然后train下是cat和dog两个文件夹,里面存的是自己的图片数据,val文件夹同train。这样我们的数据集就准备好了。
在这里插入图片描述
ImageFolder能够以目录名作为标签来对数据集做划分,下面是pytorch中文文档中关于ImageFolder的介绍:
在这里插入图片描述

#对训练集做一个变换
train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),		#对图片尺寸做一个缩放切割
    transforms.RandomHorizontalFlip(),		#水平翻转
    transforms.ToTensor(),					#转化为张量
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))	#进行归一化
])
#对测试集做变换
val_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])

train_dir = "G:/data/train"           #训练集路径
#定义数据集
train_datasets = datasets.ImageFolder(train_dir, transform=train_transforms)
#加载数据集
train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True)

val_dir = "G:/da
### 使用PyTorch实现自定义数据集上的图像分类训练 #### 定义自定义数据集类 为了处理自定义的数据集,在 PyTorch 中通常会创建一个继承 `torch.utils.data.Dataset` 类的新类。这个新类至少要重写两个方法:`__len__()` 和 `__getitem__()`. 这样可以方便地加载和预处理图片以及对应的标签。 ```python from torch.utils.data import Dataset, DataLoader import os from PIL import Image import matplotlib.pyplot as plt class CustomImageDataset(Dataset): def __init__(self, root_dir, labels_file, transform=None): self.root_dir = root_dir with open(labels_file, "r") as file: lines = file.readlines() self.image_files = [line.split()[0] for line in lines] self.labels = [int(line.split()[1]) for line in lines] self.transform = transform def __len__(self): return len(self.image_files) def __getitem__(self, idx): img_path = os.path.join(self.root_dir, self.image_files[idx]) image = Image.open(img_path).convert('RGB') label = self.labels[idx] if self.transform is not None: image = self.transform(image) sample = {'image': image, 'label': label} return sample ``` 上述代码展示了如何构建一个简单的自定义数据集类[^3]. #### 加载并可视化部分样本 通过实例化上面定义的 `CustomImageDataset`, 可以轻松获取到单个或多个样本用于调试目的: ```python train_dataset = CustomImageDataset( root_dir='./data/train', labels_file='./data/train_labels.txt' ) plt.figure(figsize=(8, 8)) for index, item in enumerate(train_dataset[:16]): ax = plt.subplot(4, 4, index + 1) ax.axis('off') ax.imshow(item['image'].permute(1, 2, 0)) # 如果图像是tensor则需转置维度顺序 ax.set_title(f'Label {item["label"]}') plt.show() ``` 这段脚本用来展示前十六张来自训练集合中的图片及其类别标签. #### 准备模型架构 对于不同的应用场景可以选择不同类型的卷积神经网络(CNN),比如 VGG 或 ResNet 。这里给出两种方式来调整这些预训练模型以便适应新的任务需求: - **微调VGG**: 修改最后一层全连接层使其输出数目等于目标分类数. ```python import torchvision.models as models model_vgg = models.vgg16(pretrained=True) num_ftrs = model_vgg.classifier[-1].in_features model_vgg.classifier[-1] = nn.Linear(num_ftrs, num_classes) ``` - **修改ResNet的最后一层** : 同样的思路应用于其他CNN结构上. ```python model_resnet = models.resnet50(pretrained=True) num_ftrs = model_resnet.fc.in_features model_resnet.fc = nn.Linear(num_ftrs, num_classes) ``` 以上两段代码分别说明了怎样针对特定的任务去定制化已有的 CNN 架构[^1][^2]. #### 训练过程设置 最后一步就是编写循环来进行迭代更新权重参数直到满足停止条件为止。这期间还需要监控验证集的表现情况从而决定何时保存最佳版本的模型文件。 ```python def train_model(model, dataloaders, criterion, optimizer, scheduler, device, num_epochs=25): best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch}/{num_epochs - 1}') # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': model.train() # Set model to training mode else: model.eval() running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) # zero the parameter gradients optimizer.zero_grad() # forward # track history if only in train with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # backward + optimize only if in training phase if phase == 'train': loss.backward() optimizer.step() # statistics running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) if phase == 'train': scheduler.step() epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print('{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) # deep copy the model if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) print() # load best model weights model.load_state_dict(best_model_wts) return model ``` 此函数实现了完整的训练流程,包括切换模式、计算损失值、反向传播梯度以及评估性能指标等操作[^4].
评论 35
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值