diff --git a/day5/requirements.txt b/day5/requirements.txt index d37edb49b..b3c12277c 100644 --- a/day5/requirements.txt +++ b/day5/requirements.txt @@ -4,4 +4,5 @@ mlflow pandas pytest great_expectations -black \ No newline at end of file +black +psutil \ No newline at end of file diff --git "a/day5/\346\274\224\347\277\2223/tests/test_model.py" "b/day5/\346\274\224\347\277\2223/tests/test_model.py" index e11a19a5c..48f4ea819 100644 --- "a/day5/\346\274\224\347\277\2223/tests/test_model.py" +++ "b/day5/\346\274\224\347\277\2223/tests/test_model.py" @@ -11,6 +11,7 @@ from sklearn.preprocessing import OneHotEncoder, StandardScaler from sklearn.compose import ColumnTransformer from sklearn.pipeline import Pipeline +import psutil # テスト用データとモデルパスを定義 DATA_PATH = os.path.join(os.path.dirname(__file__), "../data/Titanic.csv") @@ -157,7 +158,7 @@ def test_model_reproducibility(sample_data, preprocessor): steps=[ ("preprocessor", preprocessor), ("classifier", RandomForestClassifier(n_estimators=100, random_state=42)), - ] + ] ) # 学習 @@ -171,3 +172,19 @@ def test_model_reproducibility(sample_data, preprocessor): assert np.array_equal( predictions1, predictions2 ), "モデルの予測結果に再現性がありません" + +def memory_usage_amount(train_model): + model, X_test, y_test = train_model + # プロセス取得 + process = psutil.Process(os.getpid()) + + # predict前のメモリ + mem_before = process.memory_info().rss / 1024 ** 2 # MB単位 + + # 実行 + model.predict(X_test) + + # predict後のメモリ + mem_after = process.memory_info().rss / 1024 ** 2 + + assert mem_after - mem_before < 10, "推論にメモリを使いすぎです"