def predict_picture(pic_path,model_path,dataset,device):
if dataset == 'char74k':
size = 20
num_classes = 26
else:
size = 28
num_classes = 10
transform = transforms.Compose([
transforms.Scale(size),
transforms.CenterCrop((size,size)),
transforms.ToTensor(),
transforms.Normalize((0.1307,),(0.3081,))
])
img = Image.open(pic_path)
img_tensor = transform(img)
img_tensor = img_tensor.unsqueeze(0)
img_tensor = img_tensor.to(device)
model = CapsNet(num_classes=num_classes,conv_in=3)
model = model.to(device)
model.eval()
state_dict = torch.load(model_path)
model.load_state_dict(state_dict)
score = model(img_tensor)
probs = nn.functional.softmax(score,dim=-1)
max_value,index = torch.max(probs,dim=-1)
print('The picture {} is:{}'.format(pic_path.split('/')[-1],char74k_id2labels[int(index)]))
adamOptimizer.step()
epoch_loss += loss
if batch_id % 5 == 0:
logger.info('epoch:{} || batch id:{} || train loss:{}'.format(
epoch, batch_id, loss.data))
break
del data
del target
logger.info(
'epoch:{} || epoch loss:{} || best epoch:{} || best val acc: {}'.
format(epoch, epoch_loss / float(len(dataProcessor.train_loader)),
best_val_epoch, best_val_acc))
capsNet.eval()
scores, valid_preds, valid_labels = predict(capsNet, dataProcessor,
num_classes, epoch, device,
'valid', logger)
scores_test, test_preds, test_labels = predict(capsNet, dataProcessor,
num_classes, epoch, device,
'test', logger)
if scores_test >= best_val_acc:
best_val_acc = scores_test
best_val_epoch = epoch
torch.save(capsNet.state_dict(), model_path + '/model_state_dict.pkl')
elif epoch - best_val_epoch > cfg.patience:
logger.info(
"Since the val acc has not improved after {} epoch(s), "
"you have got an excellent enough model,congratulations!".format(
cfg.patience))