I did more tests to isolate the bottleneck even more. This time I changed code a little bit to navigate the times manually. Here is train function code:

def train(self, dataloader, model, loss_func, optimizer):

    train_time = time.time()
    data_time = time.time()
    total_data_time = 0
    total_pred_time = 0
    total_loss_time = 0
    total_zero_grad_time = 0
    total_backward_time = 0
    total_step_time = 0
    total_sample_time = 0

    for batch, (imgs, labels) in enumerate(dataloader):
        total_data_time += time.time() - data_time
        sample_time = time.time()

        imgs, labels = imgs.to(self.device), labels.to(self.device)

        pred_time = time.time()
        predictions = model(imgs)
        total_pred_time += time.time() - pred_time

        loss_time = time.time()
        loss = loss_func(predictions, labels)
        total_loss_time += time.time() - loss_time

        zero_grad_time = time.time()
        optimizer.zero_grad()
        total_zero_grad_time += time.time() - zero_grad_time

        backward_time = time.time()
        loss.backward()
        total_backward_time += time.time() - backward_time

        step_time = time.time()
        optimizer.step()
        total_step_time += time.time() - step_time

        total_sample_time += time.time() - sample_time
        data_time = time.time()

    print(f"Dataset: {len(dataloader.dataset)}")
    print(f"Dataloader: {len(dataloader)}")
    print(f"Train time: {time.time() - train_time} s")
    print(f"Sample time: {total_sample_time} s")
    print(f"Data time: {total_data_time} s")
    print(f"Pred time: {total_pred_time} s")
    print(f"Loss time: {total_loss_time} s")
    print(f"Zero time: {total_zero_grad_time} s")
    print(f"Backward time: {total_backward_time} s")
    print(f"Step time: {total_step_time} s")


I did two tests on ShuffleNet V2 x0.5 with batch size 142 and 566. I chose this model because the dependence on it is the most visible. Here are the results:
For batch_size=566:
Dataset: 1188
Dataloader: 3
Train time: 4.74889 s
Total Sample time: 0.20170 s
Total Data time: 4.49197 s
Total Prediction time: 0.06613 s
Total Loss time: 0.00386 s
Total ZeroGrad time: 0.00027 s
Total Backward time: 0.00483 s
Total Step time: 0.00048 s
For batch_size=142:
Dataset: 1188
Dataloader: 9
Train time: 2.70721 s
Total Sample time: 0.20925 s
Total Data time: 2.44183 s
Total Prediction time: 0.09259 s
Total Loss time: 0.00435 s
Total ZeroGrad time: 0.00082 s
Total Backward time: 0.00579 s
Total Step time: 0.00109 s
I named it data time, but I dont know for sure whats going on between one iteration and another. I suspect its time of loading single batch from dataloader, but then why profiler shows it’s only small part of time?