Skip to content

Converting

convert_model(train_sequences, val_sequences, test_sequences, parameters, dir_path, training_done)

Convert trained model to ONNX.

Parameters:

Name Type Description Default
train_sequences List[Tuple[pd.DataFrame, float]]

Training sequences.

 required
val_sequences List[Tuple[pd.DataFrame, float]]

Validation sequences.

 required
test_sequences List[Tuple[pd.DataFrame, float]]

Test sequences.

 required
parameters str

Parameters used for training.

 required
dir_path str

Directory path where the model is saved.

 required
training_done Dict[str, bool]

Flag indicating if the training is done.

 required

Returns:

Type Description
Dict[str, Any]

Dictionary of outputs from the conversion step.
Source code in make_us_rich/pipelines/converting/nodes.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def convert_model(
    train_sequences: List[Tuple[pd.DataFrame, float]],
    val_sequences: List[Tuple[pd.DataFrame, float]],
    test_sequences: List[Tuple[pd.DataFrame, float]],
    parameters: str,
    dir_path: str,
    training_done: Dict[str, bool],
) -> Dict[str, Any]:
    """
    Convert trained model to ONNX.

    Parameters
    ----------
    train_sequences: List[Tuple[pd.DataFrame, float]]
        Training sequences.
    val_sequences: List[Tuple[pd.DataFrame, float]]
        Validation sequences.
    test_sequences: List[Tuple[pd.DataFrame, float]]
        Test sequences.
    parameters: str
        Parameters used for training.
    dir_path: str
        Directory path where the model is saved.
    training_done: Dict[str, bool]
        Flag indicating if the training is done.

    Returns
    -------
    Dict[str, Any]
        Dictionary of outputs from the conversion step.
    """
    if training_done["training_done"] == True:
        model_path = [file for file in glob.glob(f"{dir_path}/*.ckpt")][0]
        model = PricePredictor.load_from_checkpoint(model_path)
        data = LSTMDataLoader(
            train_sequences=train_sequences, 
            val_sequences=val_sequences, 
            test_sequences=test_sequences, 
            train_batch_size=parameters["batch_size"], 
            val_batch_size=parameters["batch_size"],
        )
        data.setup()
        input_batch = next(iter(data.train_dataloader()))
        input_sample = input_batch[0][0].unsqueeze(0)
        path_onnx_model = f"{dir_path}/model.onnx"
        torch.onnx.export(
            model, input_sample, path_onnx_model,
            export_params=True,
            opset_version=11,
            input_names=["sequence"],
            output_names=["output"],
            dynamic_axes={
                "sequence": {0: "batch_size"},
                "output": {0: "batch_size"},
            },
        )
    return {
        "conversion_done": True,
        "model_path": model_path,
        "input_sample": input_sample,
    }

to_numpy(tensor)

Converts a tensor to numpy.

Parameters:

Name Type Description Default
tensor torch.Tensor

Tensor to be converted.

 required

Returns:

Type Description
numpy.ndarray
Source code in make_us_rich/pipelines/converting/nodes.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def to_numpy(tensor: torch.Tensor):
    """
    Converts a tensor to numpy.

    Parameters
    ----------
    tensor: torch.Tensor
        Tensor to be converted.

    Returns
    -------
    numpy.ndarray
    """
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

validate_model(dir_path, conversion_outputs)

Check if the converted model is valid.

Parameters:

Name Type Description Default
dir_path str

Directory path where the model is saved.

 required
conversion_outputs Dict[str, Any]

Dictionary of outputs from the conversion step.

 required

Returns:

Type Description
Dict[str, bool]

Flag indicating if the model is valid.
Source code in make_us_rich/pipelines/converting/nodes.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def validate_model(
    dir_path: str, 
    conversion_outputs: Dict[str, Any],
) -> Dict[str, bool]:
    """
    Check if the converted model is valid.

    Parameters
    ----------
    dir_path: str
        Directory path where the model is saved.
    conversion_outputs: Dict[str, Any]
        Dictionary of outputs from the conversion step.

    Returns
    -------
    Dict[str, bool]
        Flag indicating if the model is valid.
    """
    if conversion_outputs["conversion_done"] == True:
        path_onnx_model = f"{dir_path}/model.onnx"
        onnx_model = onnx.load(path_onnx_model)
        try:
            onnx.checker.check_model(onnx_model)
        except onnx.checker.ValidationError as e:
            raise ValueError(f"ONNX model is not valid: {e}")

        try:
            input_sample = conversion_outputs["input_sample"]
            model = PricePredictor.load_from_checkpoint(conversion_outputs["model_path"])
            model.eval()
            with torch.no_grad():
                torch_output = model(input_sample)
            ort_session = onnxruntime.InferenceSession(path_onnx_model)
            ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_sample)}
            ort_outputs = ort_session.run(None, ort_inputs)
            np.testing.assert_allclose(to_numpy(torch_output), ort_outputs[0], rtol=1e-03, atol=1e-05)
            print("🎉 ONNX model is valid. 🎉")
        except Exception as e:
            raise ValueError(f"ONNX model is not valid: {e}")
        return {"validation_done": True}
Last update: 2022-05-04