def predictAll(model_name, model_class, weight_pth, image_size, normalize):
print("[+] {0} predictAll.".format(model_name))
model = get_model(model_class)
model.load_state_dict(torch.load(weight_pth)['state_dict'])
model.eval()
print('{0} load state dict done'.format(model_name))
tta_preprocess = [preprocess(normalize, image_size), preprocess_hflip(normalize, image_size)]
tta_preprocess += make_transforms([transforms.Resize((image_size + 20, image_size + 20))],
[transforms.ToTensor(), normalize],
five_crops(image_size))
tta_preprocess += make_transforms([transforms.Resize((image_size + 20, image_size + 20))],
[HorizontalFlip(), transforms.ToTensor(), normalize],
five_crops(image_size))
print('[+] tta size: {0}'.format(len(tta_preprocess)))
data_loaders = []
for transform in tta_preprocess:
# test_dataset,_ = split_Dataset(data_dir, ratio, image_size, train_transform=transform)
test_dataset = MyDataset(test_inputs,test_labels,transform=transform)
data_loader = DataLoader(dataset=test_dataset, num_workers=16,
batch_size=BATCH_SIZE,
shuffle=False)
data_loaders.append(data_loader)
print('add transforms')
lx, px = utils.predict_tta(model, data_loaders)
data = {
'lx': lx.cpu(),
'px': px.cpu(),
}
if not os.path.exists('../feature/'+model_name):
os.makedirs('../feature/'+model_name)
torch.save(data, '../feature/'+model_name+'/all_prediction.pth')
print('{0} Predict Done'.format(model_name))
def predict(model_name, model_class, weight_pth, image_size, normalize):
print(f'[+] predict {model_name}')
model = get_model(model_class)
model.load_state_dict(torch.load(weight_pth))
model.eval()
tta_preprocess = [
preprocess(normalize, image_size),
preprocess_hflip(normalize, image_size)
]
tta_preprocess += make_transforms(
[transforms.Resize((image_size + 20, image_size + 20))],
[transforms.ToTensor(), normalize], five_crops(image_size))
tta_preprocess += make_transforms(
[transforms.Resize((image_size + 20, image_size + 20))],
[HorizontalFlip(), transforms.ToTensor(), normalize],
five_crops(image_size))
print(f'[+] tta size: {len(tta_preprocess)}')
data_loaders = []
for transform in tta_preprocess:
test_dataset = FurnitureDataset('test', transform=transform)
data_loader = DataLoader(dataset=test_dataset,
num_workers=1,
batch_size=BATCH_SIZE,
shuffle=False)
data_loaders.append(data_loader)
lx, px = utils.predict_tta(model, data_loaders)
data = {
'lx': lx.cpu(),
'px': px.cpu(),
}
torch.save(data, f'{model_name}_test_prediction.pth')
data_loaders = []
for transform in tta_preprocess:
test_dataset = FurnitureDataset('val', transform=transform)
data_loader = DataLoader(dataset=test_dataset,
num_workers=1,
batch_size=BATCH_SIZE,
shuffle=False)
data_loaders.append(data_loader)
lx, px = utils.predict_tta(model, data_loaders)
data = {
'lx': lx.cpu(),
'px': px.cpu(),
}
torch.save(data, f'{model_name}_val_prediction.pth')
def predict_to_ensemble(model_name,
model_class,
model_state_pth,
image_size,
normalize,
nb_classes=15,
batch_size=15,
with_crops=True):
print(f'[+] predict {model_name}')
model = get_model(model_class, nb_classes, model_state_pth=model_state_pth)
model.eval()
tta_preprocess = [
preprocess(normalize, image_size),
preprocess_hflip(normalize, image_size)
]
if with_crops:
tta_preprocess += make_transforms(
[transforms.Resize((image_size + 20, image_size + 20))],
[transforms.ToTensor(), normalize], five_crops(image_size))
print(f'[+] tta size: {len(tta_preprocess)}')
data_loaders = []
for transform in tta_preprocess:
data_loader = get_data_loader('./test/',
data_transform=transform,
batch_size=batch_size)
data_loaders.append(data_loader)
lx, px = utils.predict_tta(model, data_loaders)
test_predict = {
'lx': lx.cpu(),
'px': px.cpu(),
}
torch.save(test_predict, f'{model_name}_test_prediction.pth')
data_loaders = []
for transform in tta_preprocess:
valid_dataset = ImageFolder('./data/train/', transform=transform)
data_loader = get_data_loader('./data/train/',
batch_size=batch_size,
dataset=valid_dataset)
data_loaders.append(data_loader)
lx, px = utils.predict_tta(model, data_loaders)
val_predict = {
'lx': lx.cpu(),
'px': px.cpu(),
}
torch.save(val_predict, f'{model_name}_val_prediction.pth')
return {'test': test_predict, 'val': val_predict}