Обучение модели PyTorch с помощью DataLoader происходит слишком медленно ⇐ Python

[Anonymous]

Я обучаю очень маленькую НС, используя набор данных HAM10000. Для загрузки данных я использую DataLoader, который поставляется с PyTorch:

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

class CocoDetectionWithFilenames(CocoDetection):
def __init__(self, root: str, ann_file: str, transform=None):
super().__init__(root, ann_file, transform)

def get_filename(self, idx: int) -> str:
return self.coco.loadImgs(self.ids[idx])[0]["file_name"]

def get_loaders(root: str, ann_file: str) -> tuple[CocoDetection, DataLoader, DataLoader, DataLoader]:
transform = transforms.Compose([
transforms.ToTensor()
])
dataset = CocoDetectionWithFilenames(
root=root,
ann_file=ann_file,
transform=transform
)
train_size = int(0.7 * len(dataset))
valid_size = int(0.15 * len(dataset))
test_size = len(dataset) - train_size - valid_size
train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, valid_size, test_size])
num_workers = os.cpu_count()
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=32,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
prefetch_factor=1024
)
valid_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=32,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
prefetch_factor=1024
)
test_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=32,
shuffle=False,
num_workers=num_workers,
pin_memory=True
)

return dataset, train_loader, valid_loader, test_loader

Дело в том, что когда мой цикл обучения выполняется, само обучение происходит очень быстро, но программа тратит 95% времени на промежутки между эпохами — вероятно, загружая данные:

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

def extract_bboxes(targets: list[dict]) -> list[torch.Tensor]:
bboxes = []

for target in targets:
xs, ys, widths, heights = target["bbox"]

for idx, _ in enumerate(xs):
x1, y1, width, height = xs[idx], ys[idx], widths[idx], heights[idx]
# Convert COCO format (x, y, width, height) to (x1, y1, x2, y2)
x2, y2 = x1 + width, y1 + height

bboxes.append(torch.IntTensor([x1, y1, x2, y2]))

return bboxes

num_epochs = 25
train_losses = []
val_losses = []

for epoch in range(num_epochs):
model.train()
running_loss = 0.0

for images, targets in train_loader_tqdm:
images = images.to(device)
bboxes = extract_bboxes(targets)
bboxes = torch.stack(bboxes).to(device)

optimizer.zero_grad(set_to_none=True)

outputs = model(images)
loss = criterion(outputs, bboxes)

loss.backward()
optimizer.step()

running_loss += loss.item()

epoch_train_loss = running_loss / len(train_loader)

train_losses.append(epoch_train_loss)
print(f"Epoch {epoch + 1}, Loss: {epoch_train_loss}")
model.eval()

Как видите, код цикла обучения довольно прост, ничего странного в нем не происходит.

Подробнее здесь: https://stackoverflow.com/questions/786 ... s-too-slow