Ошибка: размер тензора A (810) должен соответствовать размеру тензора B (36) в не-синглтонском измерении 3 во время вычи ⇐ Python

[Anonymous]

В настоящее время я пытаюсь обучить модель обнаружения объектов с помощью Pytorch и библиотеки Effdet. Тем не менее, я сталкиваюсь со временем выполнения во время учебного цикла при расчете потери. Сообщение об ошибке:

Код: Выделить всё

RuntimeError: The size of tensor a (810) must match the size of tensor b (36) at non-singleton dimension 3 

Ниже приведен полный код, который я пробовал:
import torch
import torchvision
import numpy as np
import json
import os
import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2
from torch.utils.data import Dataset, DataLoader
from effdet import get_efficientdet_config, EfficientDet, DetBenchTrain, DetBenchPredict
from effdet.efficientdet import HeadNet
import timm.optim
from tqdm import tqdm
import matplotlib.pyplot as plt
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
< /code>
# ----------------------
# Configuration
# ----------------------
BATCH_SIZE = 4
NUM_EPOCHS = 10
IMG_SIZE = 256
NUM_CLASSES = 4 # Change based on your dataset
MODEL_NAME = "tf_efficientdet_d3"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
< /code>
class PlantDiseaseDataset(Dataset):
def __init__(self, annotation_file, image_dir, transform=None):
self.image_dir = image_dir
self.transform = transform

# Load JSON Annotations (COCO format)
with open(annotation_file, 'r') as f:
self.annotations = json.load(f)

# Ensure proper structure for COCO annotations
self.images = self.annotations.get('images', [])
self.annotations_list = self.annotations.get('annotations', [])
self.categories = self.annotations.get('categories', [])

# Debug: Show structure of data
print(f"Found {len(self.images)} images, {len(self.annotations_list)} annotations, and {len(self.categories)} categories.")

# Raise error if no images or annotations are found
if len(self.images) == 0 or len(self.annotations_list) == 0:
raise ValueError("Dataset has no images or annotations! Check your annotation file.")

def __len__(self):
return len(self.images)

def __getitem__(self, idx):
# Get the image data
img_info = self.images[idx]
img_path = os.path.join(self.image_dir, img_info['file_name'])
image = cv2.imread(img_path)

# Convert BGR to RGB
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Convert the image to a numpy array (float32) as expected by albumentations
image = np.array(image).astype(np.float32)

# Get annotations for this image
annotations = [ann for ann in self.annotations_list if ann['image_id'] == img_info['id']]

# Extract bounding boxes and labels
boxes = []
labels = []
height, width, _ = image.shape
for ann in annotations:
x, y, w, h = ann['bbox']
# Normalize the bounding boxes
x_min = x / width
y_min = y / height
x_max = (x + w) / width
y_max = (y + h) / height
boxes.append([x_min, y_min, x_max, y_max])
labels.append(ann['category_id'])

# Convert bounding boxes and labels to tensors
boxes = torch.tensor(boxes, dtype=torch.float32)
labels = torch.tensor(labels, dtype=torch.long)

# Apply transformations
if self.transform:
transformed = self.transform(image=image, bboxes=boxes.tolist(), labels=labels.tolist())
image = transformed['image']
boxes = torch.tensor(transformed['bboxes'], dtype=torch.float32)
labels = torch.tensor(transformed['labels'], dtype=torch.long)

# Ensure the target is formatted correctly (tensor with both bbox and label)
target = {'bbox': boxes, 'cls': labels}
return {'image': image, 'bboxes': boxes, 'labels': labels, 'target': target}
< /code>
def collate_fn(batch):
# Stack images into a single tensor
images = [item['image'] for item in batch]
images = torch.stack(images) # This converts the list into a tensor

# Stack bounding boxes and labels
bboxes = [item['bboxes'] for item in batch]
bboxes = torch.stack(bboxes)

labels = [item['labels'] for item in batch]
labels = torch.stack(labels)

# Return the batched images and targets as a tuple
return images, {'bbox': bboxes, 'cls': labels}
< /code>
# Define image normalization (common values for pre-trained models like ResNet)
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]

# Transformation for training and testing
train_transforms = A.Compose([
A.Resize(IMG_SIZE, IMG_SIZE), # Resize images to a fixed size for training
A.Normalize(mean=MEAN, std=STD), # Normalize the images
ToTensorV2() # Convert the image to a tensor
], bbox_params=A.BboxParams(format='coco', label_fields=['labels'])) # Use 'coco' format for bbox

test_transforms = A.Compose([
A.Resize(IMG_SIZE, IMG_SIZE), # Resize images to the same size for testing
A.Normalize(mean=MEAN, std=STD), # Normalize the images
ToTensorV2() # Convert the image to a tensor
], bbox_params=A.BboxParams(format='coco', label_fields=['labels'])) # Same for test data
< /code>
# ----------------------
# Load Dataset
# ----------------------
train_dataset = PlantDiseaseDataset("/content/Capstone_Team_Nugget_Final-5/train/_annotations.coco.json", "/content/Capstone_Team_Nugget_Final-5/train", transform=train_transforms)
test_dataset = PlantDiseaseDataset("/content/Capstone_Team_Nugget_Final-5/test/_annotations.coco.json", "/content/Capstone_Team_Nugget_Final-5/test", transform=test_transforms)
val_dataset = PlantDiseaseDataset("/content/Capstone_Team_Nugget_Final-5/valid/_annotations.coco.json", "/content/Capstone_Team_Nugget_Final-5/valid", transform=test_transforms)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)

< /code>
# ----------------------
# Model Setup
# ----------------------
def get_efficientdet_model(num_classes):
config = get_efficientdet_config(MODEL_NAME)
config.image_size = (256, 256) # Ensure it matches the resized input
model = EfficientDet(config)
# Load pretrained weights manually
checkpoint = torch.hub.load_state_dict_from_url(config.url, map_location="cpu")
model.load_state_dict(checkpoint)

model.class_net = HeadNet(config, num_outputs=num_classes)
return DetBenchTrain(model, create_labeler=True).to(DEVICE)

model = get_efficientdet_model(NUM_CLASSES)
< /code>
# ----------------------
# Optimizer & Scheduler
# ----------------------
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
< /code>
# ----------------------
# Evaluation Function
# ----------------------
def evaluate(model, val_loader):
model.eval()
with torch.no_grad():
total_loss = 0
for images, targets in val_loader:
images = torch.stack(images).to(DEVICE)
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]

loss_dict = model(images, targets)
total_loss += loss_dict["loss"].item()

avg_loss = total_loss / len(val_loader)
print(f"Validation Loss: {avg_loss:.4f}")
model.train()
< /code>
def train_model(model, train_loader, val_loader, optimizer, scheduler, num_epochs):
model.train()
for epoch in range(num_epochs):
epoch_loss = 0.0
progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")

for images, targets in progress_bar:
# Send images and targets to the correct device
images = images.to(DEVICE)
targets = {k: v.to(DEVICE) for k, v in targets.items()}

# Forward pass
loss_dict = model(images, targets)

# Compute total loss and backpropagate
loss = loss_dict["loss"]
loss.backward()
optimizer.step()

epoch_loss += loss.item()
progress_bar.set_postfix(loss=loss.item())

scheduler.step()
print(f"Epoch {epoch+1}/{num_epochs} Loss: {epoch_loss/len(train_loader):.4f}")

# Evaluate after each epoch
evaluate(model, val_loader)
< /code>
# ----------------------
# Train the Model
# ----------------------
train_model(model, train_loader, val_loader, optimizer, scheduler, NUM_EPOCHS)
< /code>
StackTrace:
Epoch 1/10: 0%| | 0/2241 [00:01

Подробнее здесь: https://stackoverflow.com/questions/793 ... -at-non-si