debug_dataloader.py

import os
import sys
import yaml
import time
import shutil
import torch
import random
import argparse
import datetime
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models

from torch.utils import data
from tqdm import tqdm

from ptsemseg.models import get_model
from ptsemseg.loss import get_loss_function
from ptsemseg.loader import get_loader 
from ptsemseg.utils import get_logger
from ptsemseg.metrics import runningScore, averageMeter
from ptsemseg.augmentations import get_composed_augmentations
from ptsemseg.schedulers import get_scheduler
from ptsemseg.optimizers import get_optimizer

from tensorboardX import SummaryWriter


import sys
sys.path.append('/home/tejus/lane-seg-experiments/Segmentation/')
import torch
import torch.nn as nn
import numpy as np
from torchvision import datasets, models, transforms
from datasets.kitti.config import CONFIG
from datasets.kitti.kitti_loader import kittiLoader
from datasets.tusimple.tusimple_loader import tusimpleLoader
from CAN import CAN
from datasets.tusimple.augmentations import *
import torchvision
import torch.nn.functional as F
from torch.autograd import Variable
import matplotlib.pyplot as plt
from metrics import runningScore 
from datetime import datetime
import math

# Definitions

TRAIN_BATCH = 12
VAL_BATCH = 12

augmentations = Compose([RandomRotate(5), RandomHorizontallyFlip()])
train_dataset = tusimpleLoader('/home/lanenet/train_set/', split="train", augmentations=augmentations)
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=TRAIN_BATCH, shuffle=True, num_workers=TRAIN_BATCH, pin_memory=True)

val_dataset = tusimpleLoader('/home/lanenet/train_set/', split="val", augmentations=None)
valloader = torch.utils.data.DataLoader(val_dataset, batch_size=VAL_BATCH, shuffle=False, num_workers=VAL_BATCH, pin_memory=True)

# imgs_sum = 0
# labels_sum = 0
# i = 0
# for data in trainloader:
#     i += 1
#     print(str(12*i)+'/3000')
#     imgs, labels = data
#     imgs_sum += torch.sum(imgs).item()
#     labels_sum += torch.sum(labels).item()

# print(imgs_sum, labels_sum)

imgs_sum = 0
labels_sum = 0
i = 0
for data in valloader:
    i += 1
    print(str(12*i)+'/600')
    imgs, labels = data
    imgs_sum += torch.sum(imgs).item()
    labels_sum += torch.sum(labels).item()

print(imgs_sum, labels_sum)