text_plain_output_1.png
from torchvision import models, transforms, datasets

import os

import torch

torch.__version__

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

data_dir = '../input/dogscats/dogscats/dogscats/'

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
vgg_format = transforms.Compose([transforms.CenterCrop(224), transforms.ToTensor(), normalize])

dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x), vgg_format) for x in ['train', 'valid']}

load_train = torch.utils.data.DataLoader(dsets['train'], batch_size=64, shuffle=True, num_workers=6)

load_test = torch.utils.data.DataLoader(dsets['valid'], batch_size=5, shuffle=True, num_workers=6)

count = 1
for data in load_test:
    print(count, end=',')
    if count == 1:
        inputs_try, labels_try = data
    count += 1