Notebook code for training linear regression model

train_linear_regression.py

from sagemaker.sklearn.estimator import SKLearn

FRAMEWORK_VERSION = "0.23-1"
script_path = 'source/train_linear_regression.py'

sklearn_linear_regression = SKLearn(
    entry_point=script_path,
    framework_version=FRAMEWORK_VERSION,
    instance_type="ml.c4.xlarge",
    role=role,
    sagemaker_session=sagemaker_session)

# Train the estimator on S3 training data
sklearn_linear_regression.fit({"train": train_input, "validation": val_input})

# deploy the model to create a predictor
predictor_linear_regression = sklearn_linear_regression.deploy(initial_instance_count=1, instance_type='ml.t2.medium')


# compute MSE
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score

test_y_preds = predictor_linear_regression.predict(test_x.values)
rmse_linear_regression = mean_squared_error(test_y_true, test_y_preds, squared=False)
r2_linear_regression = r2_score(test_y_true, test_y_preds)
print(f"Normalized RMSE: {rmse_linear_regression/normalization_factor}")
print(f"R-Squared Score: {r2_linear_regression}")