基于Faster R-CNN的目标检测训练报错:训练模式下targets不可为None
这是一个从零开始的目标检测项目,基于Colab、PyTorch、PIL开发,使用预训练Faster R-CNN模型并自定义了检测器层,数据存储在Google Drive中。运行训练代码时触发错误:AssertionError: targets should not be none when in training mode。
原Dataset类和训练循环代码如下:
class BurgerDataset(torch.utils.data.Dataset): def __init__(self, image_folder, annotation_folder, transform=None): self.image_folder = image_folder self.annotation_folder = annotation_folder self.transform = transform self.image_names = os.listdir(image_folder) self.annotation_names = os.listdir(annotation_folder) def __len__(self): return len(self.image_names) def __getitem__(self, index): # 2 != 3 image_name = self.image_names[index] annotation_name = os.path.splitext(image_name)[0] + ".json" image_path = os.path.join(self.image_folder, image_name) annotation_path = os.path.join(self.annotation_folder, annotation_name) image = Image.open(image_path).convert("RGB") with open(annotation_path, "r") as file: data = json.load(file) annotations = data["annotations"] targets = [] labels = [] for annotation in annotations: box = annotation["bbox"] label = annotation["category_id"] targets.append(box) labels.append(label) targets = torch.FloatTensor(targets) labels = torch.LongTensor(labels) if self.transform: image = self.transform(image) return image, targets, labels
训练循环代码:
num_epochs = 10 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for epoch in range(num_epochs): model.train() train_loss = 0 for images, targets, labels in train_dataloader: images = images.to(device) targets = targets.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(images, targets)['pred'] loss = your_loss_function(outputs, targets, labels) loss.backward() optimizer.step() train_loss += loss.item() train_loss /= len(train_dataloader) model.eval() test_loss = 0 with torch.no_grad(): for images, targets in test_dataloader: images = images.to(device) targets = targets.to(device) outputs = model(images) loss = your_loss_function(outputs, targets) test_loss += loss.item() test_loss /= len(test_dataloader) print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}")
错误原因
PyTorch官方实现的Faster R-CNN在训练阶段要求输入的targets是每个样本对应一个字典,字典必须包含boxes(边界框张量)和labels(类别标签张量)两个核心键,而不是单独传入boxes和labels张量。另外,由于不同样本的目标数量不同,DataLoader默认的堆叠逻辑会报错,需要自定义collate_fn来处理批量数据。
解决方案
1. 修改Dataset的返回格式
把__getitem__的返回值从image, targets, labels改成image, target_dict,其中target_dict是包含boxes和labels的字典,同时注意将边界框格式转换为Faster R-CNN要求的[xmin, ymin, xmax, ymax](如果原标注是[x,y,w,h]格式):
class BurgerDataset(torch.utils.data.Dataset): def __init__(self, image_folder, annotation_folder, transform=None): self.image_folder = image_folder self.annotation_folder = annotation_folder self.transform = transform self.image_names = os.listdir(image_folder) self.annotation_names = os.listdir(annotation_folder) def __len__(self): return len(self.image_names) def __getitem__(self, index): image_name = self.image_names[index] annotation_name = os.path.splitext(image_name)[0] + ".json" image_path = os.path.join(self.image_folder, image_name) annotation_path = os.path.join(self.annotation_folder, annotation_name) image = Image.open(image_path).convert("RGB") with open(annotation_path, "r") as file: data = json.load(file) annotations = data["annotations"] boxes = [] labels = [] for annotation in annotations: box = annotation["bbox"] # 将[x,y,w,h]转换为[xmin, ymin, xmax, ymax] boxes.append([box[0], box[1], box[0]+box[2], box[1]+box[3]]) labels.append(annotation["category_id"]) boxes = torch.FloatTensor(boxes) labels = torch.LongTensor(labels) if self.transform: image = self.transform(image) # 构造符合要求的target字典 target = {} target["boxes"] = boxes target["labels"] = labels return image, target
2. 自定义collate_fn处理批量数据
由于每个样本的目标数量不一致,DataLoader默认的堆叠逻辑会报错,需要自定义collate_fn来处理:
def collate_fn(batch): images = [] targets = [] for img, tgt in batch: images.append(img) targets.append(tgt) # 图像张量可以堆叠,targets保持列表形式 return torch.stack(images, dim=0), targets
创建DataLoader时指定该函数:
train_dataloader = torch.utils.data.DataLoader( BurgerDataset(train_image_folder, train_anno_folder, transform=train_transform), batch_size=4, shuffle=True, collate_fn=collate_fn ) test_dataloader = torch.utils.data.DataLoader( BurgerDataset(test_image_folder, test_anno_folder, transform=test_transform), batch_size=4, shuffle=False, collate_fn=collate_fn )
3. 调整训练循环的模型调用逻辑
Faster R-CNN在训练时会自动计算损失,不需要手动调用自定义损失函数,直接传入images和targets即可获取损失值:
num_epochs = 10 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for epoch in range(num_epochs): model.train() train_loss = 0 for images, targets in train_dataloader: images = images.to(device) # 将每个target的张量移到GPU for tgt in targets: tgt["boxes"] = tgt["boxes"].to(device) tgt["labels"] = tgt["labels"].to(device) optimizer.zero_grad() # 训练模式下传入images和targets,模型返回损失字典 loss_dict = model(images, targets) # 汇总所有损失项 losses = sum(loss for loss in loss_dict.values()) losses.backward() optimizer.step() train_loss += losses.item() train_loss /= len(train_dataloader) model.eval() test_loss = 0 with torch.no_grad(): for images, targets in test_dataloader: images = images.to(device) for tgt in targets: tgt["boxes"] = tgt["boxes"].to(device) tgt["labels"] = tgt["labels"].to(device) # 测试阶段传入targets可计算验证损失 loss_dict = model(images, targets) losses = sum(loss for loss in loss_dict.values()) test_loss += losses.item() test_loss /= len(test_dataloader) print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}")
内容的提问来源于stack exchange,提问作者Krilaria

