Save and load model and artifacts¶
In this notebook I will show the different options to save and load a model, as well as some additional objects produced during training.
On a given day, you train a model...
In [1]:
Copied!
import pickle
import numpy as np
import pandas as pd
import torch
import shutil
from pytorch_widedeep.preprocessing import TabPreprocessor
from pytorch_widedeep.training import Trainer
from pytorch_widedeep.callbacks import EarlyStopping, ModelCheckpoint, LRHistory
from pytorch_widedeep.models import TabMlp, WideDeep
from pytorch_widedeep.metrics import Accuracy
from pytorch_widedeep.datasets import load_adult
from sklearn.model_selection import train_test_split
import pickle
import numpy as np
import pandas as pd
import torch
import shutil
from pytorch_widedeep.preprocessing import TabPreprocessor
from pytorch_widedeep.training import Trainer
from pytorch_widedeep.callbacks import EarlyStopping, ModelCheckpoint, LRHistory
from pytorch_widedeep.models import TabMlp, WideDeep
from pytorch_widedeep.metrics import Accuracy
from pytorch_widedeep.datasets import load_adult
from sklearn.model_selection import train_test_split
/Users/javierrodriguezzaurin/.pyenv/versions/3.10.13/envs/widedeep310/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
In [2]:
Copied!
df = load_adult(as_frame=True)
df.head()
df = load_adult(as_frame=True)
df.head()
Out[2]:
age | workclass | fnlwgt | education | educational-num | marital-status | occupation | relationship | race | gender | capital-gain | capital-loss | hours-per-week | native-country | income | |
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
0 | 25 | Private | 226802 | 11th | 7 | Never-married | Machine-op-inspct | Own-child | Black | Male | 0 | 0 | 40 | United-States | <=50K |
1 | 38 | Private | 89814 | HS-grad | 9 | Married-civ-spouse | Farming-fishing | Husband | White | Male | 0 | 0 | 50 | United-States | <=50K |
2 | 28 | Local-gov | 336951 | Assoc-acdm | 12 | Married-civ-spouse | Protective-serv | Husband | White | Male | 0 | 0 | 40 | United-States | >50K |
3 | 44 | Private | 160323 | Some-college | 10 | Married-civ-spouse | Machine-op-inspct | Husband | Black | Male | 7688 | 0 | 40 | United-States | >50K |
4 | 18 | ? | 103497 | Some-college | 10 | Never-married | ? | Own-child | White | Female | 0 | 0 | 30 | United-States | <=50K |
In [3]:
Copied!
# For convenience, we'll replace '-' with '_'
df.columns = [c.replace("-", "_") for c in df.columns]
# binary target
df["target"] = (df["income"].apply(lambda x: ">50K" in x)).astype(int)
df.drop("income", axis=1, inplace=True)
df.head()
# For convenience, we'll replace '-' with '_'
df.columns = [c.replace("-", "_") for c in df.columns]
# binary target
df["target"] = (df["income"].apply(lambda x: ">50K" in x)).astype(int)
df.drop("income", axis=1, inplace=True)
df.head()
Out[3]:
age | workclass | fnlwgt | education | educational_num | marital_status | occupation | relationship | race | gender | capital_gain | capital_loss | hours_per_week | native_country | target | |
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
0 | 25 | Private | 226802 | 11th | 7 | Never-married | Machine-op-inspct | Own-child | Black | Male | 0 | 0 | 40 | United-States | 0 |
1 | 38 | Private | 89814 | HS-grad | 9 | Married-civ-spouse | Farming-fishing | Husband | White | Male | 0 | 0 | 50 | United-States | 0 |
2 | 28 | Local-gov | 336951 | Assoc-acdm | 12 | Married-civ-spouse | Protective-serv | Husband | White | Male | 0 | 0 | 40 | United-States | 1 |
3 | 44 | Private | 160323 | Some-college | 10 | Married-civ-spouse | Machine-op-inspct | Husband | Black | Male | 7688 | 0 | 40 | United-States | 1 |
4 | 18 | ? | 103497 | Some-college | 10 | Never-married | ? | Own-child | White | Female | 0 | 0 | 30 | United-States | 0 |
In [4]:
Copied!
train, valid = train_test_split(df, test_size=0.2, stratify=df.target)
# the test data will be used lately as if it was "fresh", new data coming after some time...
valid, test = train_test_split(valid, test_size=0.5, stratify=valid.target)
train, valid = train_test_split(df, test_size=0.2, stratify=df.target)
# the test data will be used lately as if it was "fresh", new data coming after some time...
valid, test = train_test_split(valid, test_size=0.5, stratify=valid.target)
In [5]:
Copied!
print(f"train shape: {train.shape}")
print(f"valid shape: {valid.shape}")
print(f"test shape: {test.shape}")
print(f"train shape: {train.shape}")
print(f"valid shape: {valid.shape}")
print(f"test shape: {test.shape}")
train shape: (39073, 15) valid shape: (4884, 15) test shape: (4885, 15)
In [6]:
Copied!
cat_embed_cols = [
"workclass",
"education",
"marital_status",
"occupation",
"relationship",
"race",
"gender",
"capital_gain",
"capital_loss",
"native_country",
]
continuous_cols = ["age", "hours_per_week"]
cat_embed_cols = [
"workclass",
"education",
"marital_status",
"occupation",
"relationship",
"race",
"gender",
"capital_gain",
"capital_loss",
"native_country",
]
continuous_cols = ["age", "hours_per_week"]
In [7]:
Copied!
tab_preprocessor = TabPreprocessor(
embed_cols=cat_embed_cols,
continuous_cols=continuous_cols,
)
X_tab_train = tab_preprocessor.fit_transform(train)
y_train = train.target.values
X_tab_valid = tab_preprocessor.transform(valid)
y_valid = valid.target.values
tab_preprocessor = TabPreprocessor(
embed_cols=cat_embed_cols,
continuous_cols=continuous_cols,
)
X_tab_train = tab_preprocessor.fit_transform(train)
y_train = train.target.values
X_tab_valid = tab_preprocessor.transform(valid)
y_valid = valid.target.values
/Users/javierrodriguezzaurin/Projects/pytorch-widedeep/pytorch_widedeep/preprocessing/tab_preprocessor.py:358: UserWarning: Continuous columns will not be normalised warnings.warn("Continuous columns will not be normalised")
In [8]:
Copied!
tab_mlp = TabMlp(
column_idx=tab_preprocessor.column_idx,
cat_embed_input=tab_preprocessor.cat_embed_input,
cat_embed_dropout=0.1,
continuous_cols=continuous_cols,
cont_norm_layer="layernorm",
embed_continuous_method="standard",
cont_embed_dim=8,
mlp_hidden_dims=[64, 32],
mlp_dropout=0.2,
mlp_activation="leaky_relu",
)
model = WideDeep(deeptabular=tab_mlp)
tab_mlp = TabMlp(
column_idx=tab_preprocessor.column_idx,
cat_embed_input=tab_preprocessor.cat_embed_input,
cat_embed_dropout=0.1,
continuous_cols=continuous_cols,
cont_norm_layer="layernorm",
embed_continuous_method="standard",
cont_embed_dim=8,
mlp_hidden_dims=[64, 32],
mlp_dropout=0.2,
mlp_activation="leaky_relu",
)
model = WideDeep(deeptabular=tab_mlp)
In [9]:
Copied!
model
model
Out[9]:
WideDeep( (deeptabular): Sequential( (0): TabMlp( (cat_embed): DiffSizeCatEmbeddings( (embed_layers): ModuleDict( (emb_layer_workclass): Embedding(10, 5, padding_idx=0) (emb_layer_education): Embedding(17, 8, padding_idx=0) (emb_layer_marital_status): Embedding(8, 5, padding_idx=0) (emb_layer_occupation): Embedding(16, 7, padding_idx=0) (emb_layer_relationship): Embedding(7, 4, padding_idx=0) (emb_layer_race): Embedding(6, 4, padding_idx=0) (emb_layer_gender): Embedding(3, 2, padding_idx=0) (emb_layer_capital_gain): Embedding(122, 23, padding_idx=0) (emb_layer_capital_loss): Embedding(97, 21, padding_idx=0) (emb_layer_native_country): Embedding(43, 13, padding_idx=0) ) (embedding_dropout): Dropout(p=0.1, inplace=False) ) (cont_norm): LayerNorm((2,), eps=1e-05, elementwise_affine=True) (cont_embed): ContEmbeddings( INFO: [ContLinear = weight(n_cont_cols, embed_dim) + bias(n_cont_cols, embed_dim)] (linear): ContLinear(n_cont_cols=2, embed_dim=8, embed_dropout=0.0) (dropout): Dropout(p=0.0, inplace=False) ) (encoder): MLP( (mlp): Sequential( (dense_layer_0): Sequential( (0): Linear(in_features=108, out_features=64, bias=True) (1): LeakyReLU(negative_slope=0.01, inplace=True) (2): Dropout(p=0.2, inplace=False) ) (dense_layer_1): Sequential( (0): Linear(in_features=64, out_features=32, bias=True) (1): LeakyReLU(negative_slope=0.01, inplace=True) (2): Dropout(p=0.2, inplace=False) ) ) ) ) (1): Linear(in_features=32, out_features=1, bias=True) ) )
In [10]:
Copied!
early_stopping = EarlyStopping()
model_checkpoint = ModelCheckpoint(
filepath="tmp_dir/adult_tabmlp_model",
save_best_only=True,
verbose=1,
max_save=1,
)
trainer = Trainer(
model,
objective="binary",
callbacks=[early_stopping, model_checkpoint],
metrics=[Accuracy],
)
trainer.fit(
X_train={"X_tab": X_tab_train, "target": y_train},
X_val={"X_tab": X_tab_valid, "target": y_valid},
n_epochs=4,
batch_size=256,
)
early_stopping = EarlyStopping()
model_checkpoint = ModelCheckpoint(
filepath="tmp_dir/adult_tabmlp_model",
save_best_only=True,
verbose=1,
max_save=1,
)
trainer = Trainer(
model,
objective="binary",
callbacks=[early_stopping, model_checkpoint],
metrics=[Accuracy],
)
trainer.fit(
X_train={"X_tab": X_tab_train, "target": y_train},
X_val={"X_tab": X_tab_valid, "target": y_valid},
n_epochs=4,
batch_size=256,
)
epoch 1: 100%|██████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 76.25it/s, loss=0.452, metrics={'acc': 0.7867}] valid: 100%|█████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 125.36it/s, loss=0.335, metrics={'acc': 0.8532}]
Epoch 1: val_loss improved from inf to 0.33532 Saving model to tmp_dir/adult_tabmlp_model_1.p
epoch 2: 100%|██████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 76.98it/s, loss=0.355, metrics={'acc': 0.8401}] valid: 100%|█████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 106.51it/s, loss=0.303, metrics={'acc': 0.8665}]
Epoch 2: val_loss improved from 0.33532 to 0.30273 Saving model to tmp_dir/adult_tabmlp_model_2.p
epoch 3: 100%|███████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 82.71it/s, loss=0.332, metrics={'acc': 0.849}] valid: 100%|█████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 107.80it/s, loss=0.288, metrics={'acc': 0.8757}]
Epoch 3: val_loss improved from 0.30273 to 0.28791 Saving model to tmp_dir/adult_tabmlp_model_3.p
epoch 4: 100%|███████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 79.02it/s, loss=0.32, metrics={'acc': 0.8541}] valid: 100%|█████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 127.07it/s, loss=0.282, metrics={'acc': 0.8763}]
Epoch 4: val_loss improved from 0.28791 to 0.28238 Saving model to tmp_dir/adult_tabmlp_model_4.p Model weights restored to best epoch: 4
Save model: option 1¶
save (and load) a model as you woud do with any other torch model
In [11]:
Copied!
torch.save(model, "tmp_dir/model_saved_option_1.pt")
torch.save(model, "tmp_dir/model_saved_option_1.pt")
In [12]:
Copied!
torch.save(model.state_dict(), "tmp_dir/model_state_dict_saved_option_1.pt")
torch.save(model.state_dict(), "tmp_dir/model_state_dict_saved_option_1.pt")
Save model: option 2¶
use the trainer
. The trainer
will also save the training history and the learning rate history (if learning rate schedulers are used)
In [13]:
Copied!
trainer.save(path="tmp_dir/", model_filename="model_saved_option_2.pt")
trainer.save(path="tmp_dir/", model_filename="model_saved_option_2.pt")
or the state dict
In [14]:
Copied!
trainer.save(
path="tmp_dir/",
model_filename="model_state_dict_saved_option_2.pt",
save_state_dict=True,
)
trainer.save(
path="tmp_dir/",
model_filename="model_state_dict_saved_option_2.pt",
save_state_dict=True,
)
In [15]:
Copied!
%%bash
ls tmp_dir/
%%bash
ls tmp_dir/
adult_tabmlp_model_4.p
history
model_saved_option_1.pt
model_saved_option_2.pt
model_state_dict_saved_option_1.pt
model_state_dict_saved_option_2.pt
In [16]:
Copied!
%%bash
ls tmp_dir/history/
%%bash
ls tmp_dir/history/
train_eval_history.json
Note that since we have used the ModelCheckpoint
Callback, adult_tabmlp_model_2.p
is the model state dict of the model at epoch 2, i.e. same as model_state_dict_saved_option_1.p
or model_state_dict_saved_option_2.p
.
Save preprocessors and callbacks¶
...just pickle them
In [17]:
Copied!
with open("tmp_dir/tab_preproc.pkl", "wb") as dp:
pickle.dump(tab_preprocessor, dp)
with open("tmp_dir/tab_preproc.pkl", "wb") as dp:
pickle.dump(tab_preprocessor, dp)
In [18]:
Copied!
with open("tmp_dir/eary_stop.pkl", "wb") as es:
pickle.dump(early_stopping, es)
with open("tmp_dir/eary_stop.pkl", "wb") as es:
pickle.dump(early_stopping, es)
In [19]:
Copied!
%%bash
ls tmp_dir/
%%bash
ls tmp_dir/
adult_tabmlp_model_4.p
eary_stop.pkl
history
model_saved_option_1.pt
model_saved_option_2.pt
model_state_dict_saved_option_1.pt
model_state_dict_saved_option_2.pt
tab_preproc.pkl
And that is pretty much all you need to resume training or directly predict, let's see
Run New experiment: prepare new dataset, load model, and predict¶
In [20]:
Copied!
test.head()
test.head()
Out[20]:
age | workclass | fnlwgt | education | educational_num | marital_status | occupation | relationship | race | gender | capital_gain | capital_loss | hours_per_week | native_country | target | |
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
10103 | 43 | Private | 198282 | HS-grad | 9 | Married-civ-spouse | Craft-repair | Husband | White | Male | 0 | 0 | 40 | United-States | 1 |
31799 | 20 | Private | 228686 | 11th | 7 | Married-civ-spouse | Other-service | Husband | White | Male | 0 | 0 | 40 | United-States | 0 |
19971 | 26 | Private | 291968 | HS-grad | 9 | Married-civ-spouse | Transport-moving | Husband | White | Male | 0 | 0 | 44 | United-States | 0 |
3039 | 48 | Private | 175958 | Bachelors | 13 | Divorced | Prof-specialty | Not-in-family | White | Male | 0 | 0 | 30 | United-States | 0 |
20725 | 18 | Private | 232024 | 11th | 7 | Never-married | Machine-op-inspct | Own-child | White | Male | 0 | 0 | 55 | United-States | 0 |
In [21]:
Copied!
with open("tmp_dir/tab_preproc.pkl", "rb") as tp:
tab_preprocessor_new = pickle.load(tp)
with open("tmp_dir/tab_preproc.pkl", "rb") as tp:
tab_preprocessor_new = pickle.load(tp)
In [22]:
Copied!
X_test_tab = tab_preprocessor_new.transform(test)
y_test = test.target
X_test_tab = tab_preprocessor_new.transform(test)
y_test = test.target
In [23]:
Copied!
tab_mlp_new = TabMlp(
column_idx=tab_preprocessor.column_idx,
cat_embed_input=tab_preprocessor.cat_embed_input,
cat_embed_dropout=0.1,
continuous_cols=continuous_cols,
cont_norm_layer="layernorm",
embed_continuous_method="standard",
cont_embed_dim=8,
mlp_hidden_dims=[64, 32],
mlp_dropout=0.2,
mlp_activation="leaky_relu",
)
new_model = WideDeep(deeptabular=tab_mlp)
tab_mlp_new = TabMlp(
column_idx=tab_preprocessor.column_idx,
cat_embed_input=tab_preprocessor.cat_embed_input,
cat_embed_dropout=0.1,
continuous_cols=continuous_cols,
cont_norm_layer="layernorm",
embed_continuous_method="standard",
cont_embed_dim=8,
mlp_hidden_dims=[64, 32],
mlp_dropout=0.2,
mlp_activation="leaky_relu",
)
new_model = WideDeep(deeptabular=tab_mlp)
In [24]:
Copied!
new_model.load_state_dict(torch.load("tmp_dir/model_state_dict_saved_option_2.pt"))
new_model.load_state_dict(torch.load("tmp_dir/model_state_dict_saved_option_2.pt"))
Out[24]:
<All keys matched successfully>
In [25]:
Copied!
trainer = Trainer(
model,
objective="binary",
)
trainer = Trainer(
model,
objective="binary",
)
In [26]:
Copied!
preds = trainer.predict(X_tab=X_test_tab, batch_size=32)
preds = trainer.predict(X_tab=X_test_tab, batch_size=32)
predict: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:00<00:00, 309.83it/s]
In [27]:
Copied!
from sklearn.metrics import accuracy_score
from sklearn.metrics import accuracy_score
In [28]:
Copied!
accuracy_score(y_test, preds)
accuracy_score(y_test, preds)
Out[28]:
0.8595701125895598
In [29]:
Copied!
shutil.rmtree("tmp_dir/")
shutil.rmtree("tmp_dir/")