wehub-resource-sync
56 lines
1.3 KiB
Python
56 lines
1.3 KiB
Python
import argparse
|
|
|
|
import numpy as np
|
|
import statsmodels.api as sm
|
|
from sklearn.metrics import mean_squared_error
|
|
|
|
import mlflow
|
|
import mlflow.statsmodels
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Statsmodels example")
|
|
parser.add_argument(
|
|
"--inverse-method",
|
|
type=str,
|
|
default="pinv",
|
|
help="Can be 'pinv', or 'qr'. 'pinv' uses the Moore-Penrose pseudoinverse "
|
|
"to solve the least squares problem. 'qr' uses the QR factorization. "
|
|
"(default: 'pinv')",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
# parse command-line arguments
|
|
args = parse_args()
|
|
|
|
# prepare train and test data
|
|
# Ordinary Least Squares (OLS)
|
|
np.random.seed(9876789)
|
|
nsamples = 100
|
|
x = np.linspace(0, 10, 100)
|
|
X = np.column_stack((x, x**2))
|
|
beta = np.array([1, 0.1, 10])
|
|
e = np.random.normal(size=nsamples)
|
|
X = sm.add_constant(X)
|
|
y = np.dot(X, beta) + e
|
|
|
|
# enable auto logging
|
|
mlflow.statsmodels.autolog()
|
|
|
|
with mlflow.start_run():
|
|
ols = sm.OLS(y, X)
|
|
model = ols.fit(method=args.inverse_method)
|
|
|
|
# evaluate model
|
|
y_pred = model.predict(X)
|
|
mse = mean_squared_error(y, y_pred)
|
|
|
|
# log metrics
|
|
mlflow.log_metrics({"mse": mse})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|