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))