Last active
June 11, 2019 20:24
-
-
Save abhishekkrthakur/3346a3e1f47bc8a11b73078d4eca54e9 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import torch | |
from torchvision import transforms | |
# define some re-usable stuff | |
IMAGE_SIZE = 224 | |
NUM_CLASSES = 1103 | |
BATCH_SIZE = 32 | |
device = torch.device("cuda:0") | |
IMG_MEAN = model_ft.mean | |
IMG_STD = model_ft.std | |
# make some augmentations on training data | |
train_transform = transforms.Compose([ | |
transforms.RandomResizedCrop((IMAGE_SIZE, IMAGE_SIZE)), | |
transforms.RandomHorizontalFlip(), | |
transforms.ToTensor(), | |
transforms.Normalize(IMG_MEAN, IMG_STD) | |
]) | |
# use the collections dataset class we created earlier | |
train_dataset = CollectionsDataset(csv_file='../input/folds.csv', | |
root_dir='../input/train/', | |
num_classes=NUM_CLASSES, | |
transform=train_transform) | |
# create the pytorch data loader | |
train_dataset_loader = torch.utils.data.DataLoader(train_dataset, | |
batch_size=BATCH_SIZE, | |
shuffle=True, | |
num_workers=4) | |
# push model to device | |
model_ft = model_ft.to(device) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment