baselines.py 文件源码

def fit_lr(train_X, train_y, test_X):
    """
    Use linear regression to predict.

    :param train_X:
    :param train_y:
    :param test_X:
    :return:
    """
    lr = LinearRegression()
    lr.fit(train_X, train_y)
    yhat_train = lr.predict(train_X)
    yhat_test = lr.predict(test_X)
    model = "LR int %.2f coefs %s" % (lr.intercept_, pprint(lr.coef_))

    return model, yhat_train, yhat_test