[AICE PROFESSIONAL 준비] 뇌졸증 예측

캐글의 [Stroke Prediction Dataset] 을 이용하여 뇌졸증 예측 모델을 만들어봤다.

 

https://www.kaggle.com/datasets/fedesoriano/stroke-prediction-dataset

Stroke Prediction Dataset

 

데이터 로드 및 확인

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings('ignore', category=UserWarning)

df = pd.read_csv("./data/healthcare-dataset-stroke-data.csv")
df.head()

 

 

df.info()

 

df.describe()

 

df.isnull().sum()

 

bmi 에 빈값이 있는데 전체 데이터 5,110건 중 201건이라서 적지 않은 비율이지만 삭제 하기로 결정

df.dropna(inplace=True)

 

 

데이터 시각화

numerical_col = ["age", "avg_glucose_level", "bmi"] # 수치형 데이터
categorical_cols = ["gender", "hypertension", "heart_disease", "ever_married", "work_type", "Residence_type", "smoking_status", "stroke"] # 범주형 데이터

# 수치형 데이터 시각화
fig, axes = plt.subplots(1, 3, figsize=(13, 5))

for i, col_name in enumerate(numerical_col):
    ax = axes[i]

    # 히스토그램 그리기
    ax.hist(df[col_name], bins=10, color="C"+str(i), alpha=0.7, edgecolor="black")

    # 개별 차트 제목 및 라벨 설정
    ax.set_title(f"Histogram: {col_name}")
    ax.set_xlabel("Values")
    ax.set_ylabel("Frequency")

    # 격자선 추가
    ax.grid(axis="y", alpha=0.5)

plt.tight_layout()
plt.show()

 

# 범주형 데이터 시각화
fig, axes = plt.subplots(2, 4, figsize=(13, 10))
axes_flat = axes.flatten()

for i, col_name in enumerate(categorical_cols):
    ax = axes_flat[i]

    # 막대그래프를 그리기 위해서는 먼저 데이터 집계가 필요함 (항목별로 몇 개씩 있는지 카운트)
    data_counts = df[col_name].value_counts().sort_index()

    ax.bar(data_counts.index.astype(str), data_counts.values, color="C"+str(i), alpha=0.7, edgecolor="black")
    ax.set_title(f"Bar Chart: {col_name}")
    ax.set_xlabel("Categories")
    ax.set_ylabel("Count")
    ax.tick_params(axis="x", rotation=45) # x축 글자가 겹치므로 45도 회전

plt.tight_layout()
plt.show()

 

df["gender"].value_counts()

 

 

데이터 전처리

1. id는 의미 없는 데이터이므로 삭제

2. 범주형 데이터

    - gender의 Ohter값은 1건만 존재하는데, 모델이 학습하기 어려운 양의 데이터이고, 혼동이 생길 수 있으므로 삭제

    - gender, ever_married, work_type, Residence_type, smoking_status는 원핫인코딩

3. 수치형 데이터

    - avg_glucose_level, bmi는 분포가 한쪽으로 쏠려 있고 이상치가 있으므로 이상치에 강한 RobustScaler 적용

    - age는 표준화를 적용하여 평균 0, 분산 1로 맞춤

4. 타겟 데이터 불균형

    - stroke의 1(정상)값이 압도적으로 많기 때문에 SMOTE를 활용하여 비율 맞추기

 

# id 필드 삭제
df_drop = df.drop(columns=["id"])

 

# Ohter 성별은 학습이 어렵기 때문에 삭제
df_drop.drop(index=(df[df["gender"]=="Other"].index), inplace=True)

 

# 원핫인코딩
df_encoding = pd.get_dummies(df_drop, columns=["gender","ever_married", "work_type","Residence_type","smoking_status"], drop_first=True)
df_encoding

 

df_encoding.info()

 

from sklearn.model_selection import train_test_split
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, RobustScaler

robust_cols = ["avg_glucose_level", "bmi"]
standard_cols = ["age"]

X = df_encoding.drop("stroke", axis=1)
y = df_encoding["stroke"]

# 데이터 분리
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# 전처리 파이프라인 구성
# ColumnTransformer로 컬럼에 맞는 스케일러를 매핑
preprocessor = ColumnTransformer(
    transformers=[
        ("robust", RobustScaler(), robust_cols),
        ("std", StandardScaler(), standard_cols)
    ],
    remainder="passthrough" # 스케일링 하지 않을 나머지 컬럼은 그대로 둔다
)

 

 

하이퍼파라미터 튜닝 및 모델 학습

from sklearn.ensemble import RandomForestClassifier
from imblearn.pipeline import Pipeline as ImbPipeline
from imblearn.over_sampling import SMOTE
from sklearn.model_selection import GridSearchCV

# RandomForestClassifier 전용 파이프라인 구성
pipeline_rf = ImbPipeline(
    [
        ("preprocessor", preprocessor),
        ("smote", SMOTE(random_state=42)),
        ("classifier", RandomForestClassifier(random_state=42))
    ]
)

# RandomForestClassifier 하이퍼파라미터 그리드 설정
param_grid_rf = {
    "classifier__n_estimators": [100, 200, 300],
    "classifier__max_depth": [10, 20, None],
    "classifier__min_samples_split": [2, 5, 10],
    "classifier__min_samples_leaf": [1, 2, 4]
    }

# GridSearchCV 설정
grid_search_rf = GridSearchCV(
    estimator=pipeline_rf,
    param_grid=param_grid_rf,
    cv=5,
    scoring="recall",
    n_jobs=-1,
    verbose=2
)

# 원본 데이터로 학습 시작
print("===== RandomForestClassifier GridSearchCV 학습 시작 =====")
grid_search_rf.fit(X_train, y_train)
print("최적의 파라미터: ", grid_search_rf.best_params_)

# 최적 모델 추출 및 예측
best_rf_model = grid_search_rf.best_estimator_

 

from sklearn.metrics import classification_report

# 테스트 데이터로 예측 수행
y_pred_rf = best_rf_model.predict(X_test)


# 성능 지표
# 뇌졸증 환자를 예측하는 recall 점수는 5% 로 매우 낮음
print(classification_report(y_test, y_pred_rf, target_names=["Normal[0]", "Stroke(1)"]))

 

뇌졸증 환자를 예측해야하는 모델이므로 accuracy 보다는 recall 이나 f1-score가 중요한데... 맞출 확률이 29%라는 처참한 스코어가 나왔다. 다른 모델을 사용하기로 결정했다.

 

 

XGBoost 모델

from xgboost import XGBClassifier
from imblearn.pipeline import Pipeline as ImbPipeline
from imblearn.over_sampling import SMOTE
from sklearn.model_selection import GridSearchCV

# XGBoost 전용 파이프라인 구성
pipeline_xgb = ImbPipeline(
    [
        ("preprocessor", preprocessor),
        ("smote", SMOTE(random_state=42)),
        ("classifier", XGBClassifier(random_state=42, eval_metric="logloss"))
    ]
)

# XGBoost 하이퍼파라미터 그리드 설정
param_grid_xgb = {
    "classifier__n_estimators": [100, 200],
    "classifier__max_depth": [3, 5, 7],
    "classifier__learning_rate": [0.01, 0.1, 0.2]
}

# GridSearchCV 설정
grid_search_xgb = GridSearchCV(
    estimator=pipeline_xgb,
    param_grid=param_grid_xgb,
    cv=5,
    scoring="recall",
    n_jobs=-1,
    verbose=2
)

# 원본 데이터로 학습 시작
print("===== XGBoost GridSearchCV 학습 시작 =====")
grid_search_xgb.fit(X_train, y_train)
print("최적의 파라미터: ", grid_search_xgb.best_params_)

# 최적 모델 추출 및 예측
best_xgb_model = grid_search_xgb.best_estimator_

 

y_pred_xgb = best_xgb_model.predict(X_test)

# 성능 지표
print(classification_report(y_test, y_pred_xgb, target_names=["Normal[0]", "Stroke(1)"]))

 

뇌졸중 환자를 놓치지 않고 예측한 recall 점수가 62%로 RandomForestClassifer 보다 훨씬 좋은 성능을 보였다.

 

 

LightGBM 모델

from lightgbm import LGBMClassifier

# LightGBM 전용 파이프라인 구성
pipeline_lgbm = ImbPipeline([
    ('preprocessor', preprocessor),
    ('smote', SMOTE(random_state=42)),
    ('classifier', LGBMClassifier(random_state=42, verbose=-1)) # 로그 출력 최소화
])

# LightGBM 하이퍼파라미터 그리드 설정
param_grid_lgbm = {
    'classifier__n_estimators': [100, 200],
    'classifier__max_depth': [3, 5, -1], # -1은 제한 없음
    'classifier__num_leaves': [15, 31, 63],
    'classifier__learning_rate': [0.01, 0.1]
}

# GridSearchCV 설정
grid_search_lgbm = GridSearchCV(
    estimator=pipeline_lgbm,
    param_grid=param_grid_lgbm,
    cv=5,
    scoring="recall",
    n_jobs=-1,
    verbose=2
)

# 원본 데이터로 학습 시작
print("===== LightGBM GridSearchCV 시작 =====")
grid_search_lgbm.fit(X_train, y_train)
print("최적의 파라미터: ", grid_search_lgbm.best_params_)

# 최적 모델 추출 및 예측
best_lgbm_model = grid_search_lgbm.best_estimator_

 

y_pred_lgbm = best_lgbm_model.predict(X_test)

# 성능 지표
print(classification_report(y_test, y_pred_lgbm, target_names=["Normal[0]", "Stroke(1)"]))

 

XGBoost 모델과 동일하게 recall은 62%가 나왔다. 만족할만한 수치는 아니지만 RandomForestClassifier 모델에 비하면 훨씬 좋은 성능을 보인 XGBoost와 LightGBM 모델이다. SMOTE를 사용해서 타겟 데이터의 불균형을 해소하려 한 것이 오히려 과적합 문제를 일으키지 않았나 싶다.

 

 

모델 저장

# 데이터 양이 적을 때는 보통 LightGBM 보다는 XGBoost 를 사용하므로, XGBoost 모델을 저장
import joblib

joblib.dump(best_xgb_model, "./my_xgb_model.pkl")

 

 

가상으로 만든 데이터를 이용하여 예측하기

# stroke 값이 없는 가상의 데이터를 만들어서 예측
df = pd.read_csv("./data/real-data.csv")
df.head()

 

df.info()

 

# 학습 데이터와 동일하게 전처리 진행
# 단 bmi값이 없는 데이터를 지우지 않고 평균값으로 대체하여 진행
df_fill = df.copy()
df_fill["bmi"] = df_fill["bmi"].fillna(df_fill["bmi"].mean())

# id 필드 삭제
df_drop = df_fill.drop(columns=["id"])

# Ohter 성별은 학습이 어렵기 때문에 삭제
df_drop.drop(index=(df[df["gender"]=="Other"].index), inplace=True)

# 원핫인코딩
df_encoding = pd.get_dummies(df_drop, columns=["gender","ever_married", "work_type","Residence_type","smoking_status"], drop_first=True)

df_encoding.info()

 

y_pred = best_xgb_model.predict(df_encoding)
y_pred

 

# 예측 결과를 원본 데이터 마지막에 붙여넣고 데이터 저장
df["stroke"] = y_pred
df.to_csv("./result.csv")

 

 

파이토치 딥러닝으로 전환

import pandas as pd
import numpy as np
import warnings
from sklearn.model_selection import train_test_split
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, RobustScaler
from imblearn.over_sampling import SMOTE

warnings.filterwarnings('ignore', category=UserWarning)

df = pd.read_csv("./data/healthcare-dataset-stroke-data.csv")

# 머신러닝에서 했던 전처리 과정과 동일하게 진행
df.dropna(inplace=True)

# id 필드 삭제
df_drop = df.drop(columns=["id"])

# Ohter 성별은 학습이 어렵기 때문에 삭제
df_drop.drop(index=(df[df["gender"]=="Other"].index), inplace=True)

# 원핫인코딩
df_encoding = pd.get_dummies(df_drop, columns=["gender","ever_married", "work_type","Residence_type","smoking_status"], drop_first=True)

robust_cols = ["avg_glucose_level", "bmi"]
standard_cols = ["age"]

X = df_encoding.drop("stroke", axis=1)
y = df_encoding["stroke"]

# 데이터 분리
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# 전처리 파이프라인 구성
# ColumnTransformer로 컬럼에 맞는 스케일러를 매핑
preprocessor = ColumnTransformer(
    transformers=[
        ("robust", RobustScaler(), robust_cols),
        ("std", StandardScaler(), standard_cols)
    ],
    remainder="passthrough" # 스케일링 하지 않을 나머지 컬럼은 그대로 둔다
)

 

# 학습데이터와 테스트데이터 전처리
X_train_processed = preprocessor.fit_transform(X_train)
X_test_processed = preprocessor.transform(X_test)

 

import torch

# 데이터 타입 강제 변환
if hasattr(X_train_processed, "toarray"):
    X_train_processed = X_train_processed.toarray()
if hasattr(X_test_processed, "toarray"):
    X_test_processed = X_test_processed.toarray()

X_train_np = np.asarray(X_train_processed, dtype=np.float32)
y_train_np = np.asarray(y_train, dtype=np.float32)

X_test_np = np.asarray(X_test_processed, dtype=np.float32)
y_test_np = np.asarray(y_test, dtype=np.float32)

# numpy 배열을 PyTorch Tensor 로 변환
# float32 타입으로 맞추고, y는 2차원 형태(N, 1)로 변경
X_train_tensor = torch.from_numpy(X_train_np)
y_train_tensor = torch.from_numpy(y_train_np).unsqueeze(1)

X_test_tensor = torch.from_numpy(X_test_np)
y_test_tensor = torch.from_numpy(y_test_np).unsqueeze(1)

 

from torch.utils.data import DataLoader, TensorDataset

# DataLoader 생성
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

 

import torch.nn as nn

# PyTorch 신경망 모델 정의
class StrokeNN(nn.Module):
    def __init__(self, input_dim):
        super(StrokeNN, self).__init__()

        # 입력층 -> 은닉층1
        self.layer1 = nn.Linear(input_dim, 64)
        self.relu = nn.ReLU()
        self.dropout1 = nn.Dropout(0.2)

        # 은닉층1 -> 은닉층2
        self.layer2 = nn.Linear(64, 32)
        self.dropout2 = nn.Dropout(0.1)

        # 은닉층2 -> 출력층 (이진 분류이므로 출력 노드는 1개)
        self.output_layer = nn.Linear(32, 1)

    def forward(self, x):
        x = self.relu(self.layer1(x))
        x = self.dropout1(x)
        x = self.relu(self.layer2(x))
        x = self.dropout2(x)
        x = self.output_layer(x)

        return x

 

import torch.optim as optim

# 학습 설정 및 루프 구성
input_size = X_train_tensor.shape[1]
model = StrokeNN(input_dim = input_size)

# pos_weight 계산 및 적용
# 정상 클래스 수 / 뇌졸증 클래스 수
num_negatives = (y_train == 0).sum()
num_positives = (y_train == 1).sum()
weight_ratio = num_negatives / num_positives

# 텐서로 변환하여 Loss 함수에 주입
pos_weight_tensor = torch.tensor([weight_ratio], dtype=torch.float32)
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight_tensor)

optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 100

print("===== PyTorch 모델 학습 시작 =====")
for epoch in range(epochs):
    model.train()
    epoch_loss = 0

    for batch_X, batch_y in train_loader:
        # 기울기 초기화
        optimizer.zero_grad()

        # 순전파 (예측)
        predictions = model(batch_X)

        # 손실 계산
        loss = criterion(predictions, batch_y)

        # 역전파 (기울기 계산)
        loss.backward()

        # 가중치 업데이트
        optimizer.step()

        epoch_loss += loss.item()

    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {epoch_loss/len(train_loader):.4f}")

print("===== PyTorch 모델 학습 종료 =====")

 

from sklearn.metrics import classification_report

# 테스트 데이터 평가
model.eval()
with torch.no_grad():
    # 예측값 도출
    y_pred_logits = model(X_test_tensor)

    # sigmoid를 통과시켜 0 ~ 1 사이의 확률값으로 변환
    y_pred_probs = torch.sigmoid(y_pred_logits)

    # 임계값을 0.3으로 설정하여 클래스 결정
    y_pred_classes = (y_pred_probs >= 0.3).float()

print("\n[PyTroch 신경망 모델 성능 지표]")
print(classification_report(y_test_tensor.numpy(), y_pred_classes.numpy(), target_names = ["Normal(0)", "Stroke(1)"]))

 

recall 은 64%로 머신러닝보다 소폭 증가했으나, 큰 차이는 없다.