transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
# Data-loader of testing set
transform_val = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
train_dataset = DatasetFolder(DATA_INFO.TRAIN_DIR, transform_train,
DATA_INFO.NUM_CLASSES, "train",
opt.MODEL.IMAGE_SIZE)
val_dataset = DatasetFolder(DATA_INFO.VAL_DIR, transform_val,
DATA_INFO.NUM_CLASSES, "val", opt.MODEL.IMAGE_SIZE)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.TRAIN.BATCH_SIZE,
shuffle=opt.TRAIN.SHUFFLE,
num_workers=opt.TRAIN.WORKERS)
test_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.TRAIN.BATCH_SIZE,
shuffle=False,
num_workers=opt.TRAIN.WORKERS)
# create model
logger.info('\n\nOptions:')
logger.info(pprint.pformat(opt))
DATA_INFO = cfg.DATASET
# Data-loader of testing set
transform_test = transforms.Compose([
# transforms.Resize((opt.MODEL.IMAGE_SIZE)),
# transforms.CenterCrop(opt.MODEL.INPUT_SIZE),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
test_dataset = DatasetFolder(sys.argv[3], transform_test,
DATA_INFO.NUM_CLASSES, "test",
opt.MODEL.INPUT_SIZE)
logger.info(f'{len(test_dataset)} images are found for test')
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=opt.TEST.BATCH_SIZE,
shuffle=False,
num_workers=opt.TEST.WORKERS)
last_checkpoint = torch.load(opt.TEST.CHECKPOINT)
opt.MODEL.ARCH = last_checkpoint['arch']
# create model
logger.info("using pre-trained model MobileNet")
model = MobileNetV2()