一、Win10系统下FasterRCNN训练自己的数据集
在Win10系统下训练自己的数据集,需要首先安装以下环境:
- CUDA 11.0
- CUDNN 8.0.4
- anaconda
- pytorch==1.4.0
- torchvision==0.5.0
接下来,需要用到COCO API,在anaconda的环境下,使用以下命令进行安装:
pip install cython
pip install -U setuptools
pip install pycocotools
接下来,需要自己标注数据集,并将其转换为COCO格式的JSON文件。
最后,使用FasterRCNN进行训练,具体示例代码如下:
# 导入必要的包
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torchvision.models as models
import torchvision.transforms as T
from engine import train_one_epoch, evaluate
import utils
# 加载数据集
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
train_dataset = datasets.ImageFolder(root='path/to/train/folder', transform=data_transform)
train_data_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
# 加载模型
model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 2 # 包括背景和目标两个类别
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# 定义优化器和学习率
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
# 训练模型
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
utils.train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=10)
lr_scheduler.step()
二、FasterRCNN训练COCO数据集
使用预训练的FasterRCNN模型在COCO数据集上进行微调,可以通过以下代码实现:
import torchvision.transforms as T
from torchvision.datasets import CocoDetection
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torch.utils.data import DataLoader
# 获取COCO数据集
def get_coco_dataset(root, image_set, transforms):
coco = CocoDetection(root, image_set=image_set, transform=transforms)
return coco
# 获取微调模型
def get_fine_tuning_model(num_classes):
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
# 加载训练集和测试集
data_transform = T.Compose([
T.Resize(800),
T.ToTensor()
])
train_dataset = get_coco_dataset('path/to/COCO', 'train2017', data_transform)
train_data_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=4)
test_dataset = get_coco_dataset('path/to/COCO', 'val2017', data_transform)
test_data_loader = DataLoader(test_dataset, batch_size=2, shuffle=False, num_workers=4)
# 定义模型并进行训练
num_classes = 91 # 包括背景
fine_tuning_model = get_fine_tuning_model(num_classes)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
fine_tuning_model.to(device)
params = [p for p in fine_tuning_model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
num_epochs = 10 # 训练轮数
for epoch in range(num_epochs):
train_one_epoch(fine_tuning_model, optimizer, train_data_loader, device, epoch, print_freq=10)
lr_scheduler.step()
# 在测试集上进行评估
evaluate(fine_tuning_model, test_data_loader, device=device)
三、如何训练自己的数据集
如果要训练自己的数据集,需要进行以下步骤:
- 准备图片数据集,需要包含正样本和负样本。
- 对数据集进行标注,标注包括目标的位置和类别。
- 将标注后的数据集转换为COCO格式的JSON文件。
- 按照上述代码示例,使用FasterRCNN进行训练。