Почему обучение PyTorch RetinaNet ResNet50 FPN V2 в Google Colab с графическим процессором T4 происходит так медленно? ⇐ Python

[Anonymous]

Я прекрасно настраиваю модель сетчатки Resnet50 FPN V2 с использованием pytorch с retinanet_resnet50_fpn_v2_weights.default. Мой набор данных имеет около 13 050 образцов обучения, и я использую Google Colab с графическим процессором T4. Тем не менее, процесс обучения очень медленный, занимающий более 25 минут за эпоху, даже с размером партии 8. < /P>
Моя настройка:
Модель: retinanet_resnet50_fpn_v2 с предварительно проведенными весами. BR /> Размер набора данных: 13 050 изображений для обучения.
Размер партии: 16.
GPU: NVIDIA T4 (15 ГБ VRAM) на COLAB.
Следующий код: < /p>

• Coco Transformation: < /li>
  < /ul>
  # Define transformations
  class CocoTransform:
  def __call__(self, image, target):
  image = F.to_tensor(image) # Convert PIL image to tensor
  return image, target
  
  
  Загрузчик данных:

# Dataset class
def get_coco_dataset(img_dir, ann_file):
return CocoDetection(
root=img_dir,
annFile=ann_file,
transforms=CocoTransform()
)

# Load datasets
train_dataset = get_coco_dataset(
img_dir="/content/train",
ann_file="/content/train/_annotations.coco.json"
)

val_dataset = get_coco_dataset(
img_dir="/content/valid",
ann_file="/content/valid/_annotations.coco.json"
)

train_loader = DataLoader(
train_dataset,
batch_size=16,
shuffle=True,
collate_fn=lambda x: tuple(zip(*x)),
num_workers=2, # Adjust based on your CPU cores (try 2, 4, 8, etc.)
pin_memory=True, # Recommended if training on GPU
)
val_loader = DataLoader(
val_dataset,
batch_size=8,
shuffle=False,
collate_fn=lambda x: tuple(zip(*x)),
num_workers=2,
pin_memory=True,
)
< /code>
-Загрузка модели: (я хочу заморозить основу и тренировать только головку классификации) < /p>
def get_retinanet_model(num_classes):
weights = RetinaNet_ResNet50_FPN_V2_Weights.DEFAULT
trainable_backbone_layers=0
model = retinanet_resnet50_fpn_v2(weights=weights,trainable_backbone_layers=trainable_backbone_layers)

# Get the number of input channels for the feature maps
in_channels = model.head.classification_head.conv[0][0].in_channels # Should be 256 by default

# Retrieve the number of anchors per location
num_anchors = model.head.classification_head.num_anchors

# Replace the classification head to accommodate the new number of classes
model.head.classification_head = RetinaNetClassificationHead(
in_channels=in_channels,
num_anchors=num_anchors,
num_classes=num_classes
)

# Replace the regression head
model.head.regression_head = RetinaNetRegressionHead(
in_channels=in_channels,
num_anchors=num_anchors
)

return model

-Инициализируйте модель
# Initialize the model
num_classes = 4 # Background + chair, human, table
model = get_retinanet_model(num_classes)
# Move model to GPU if available
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)

# Define optimizer and learning rate scheduler
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

-обучить функцию одной эпохи:
def train_one_epoch(model, optimizer, data_loader, device, epoch):
model.train()
for images, targets in data_loader:
# Move images to the device
images = [img.to(device) for img in images]

# Validate and process targets
processed_targets = []
valid_images = []
for i, target in enumerate(targets):
boxes = []
labels = []
for obj in target:
# Extract bbox
bbox = obj["bbox"] # Format: [x, y, width, height]
x, y, w, h = bbox

# Ensure the width and height are positive
if w > 0 and h > 0:
boxes.append([x, y, x + w, y + h]) # Convert to [x_min, y_min, x_max, y_max]
labels.append(obj["category_id"])

# Only process if there are valid boxes
if boxes:
processed_target = {
"boxes": torch.tensor(boxes, dtype=torch.float32).to(device),
"labels": torch.tensor(labels, dtype=torch.int64).to(device),
}
processed_targets.append(processed_target)
valid_images.append(images) # Add only valid images

# Skip iteration if no valid targets
if not processed_targets:
continue

# Ensure images and targets are aligned
images = valid_images

# Forward pass
loss_dict = model(images, processed_targets)
losses = sum(loss for loss in loss_dict.values())

# Backpropagation
optimizer.zero_grad()
losses.backward()
optimizer.step()

print(f"Epoch [{epoch}] Loss: {losses.item():.4f}")
< /code>
-Training Loop < /p>
# Training loop
num_epochs = 20
for epoch in range(num_epochs):
train_one_epoch(model, optimizer, train_loader, device, epoch)
lr_scheduler.step()
# Save the model's state dictionary after every epoch
model_path = f"retinanet_resnet50_epoch_{epoch + 1}.pth"
torch.save(model.state_dict(), model_path)
print(f"Model saved: {model_path}")


Подробнее здесь: https://stackoverflow.com/questions/793 ... lab-with-a