【深度学习实验】—— 基于 LSTM 的年度医疗费用回归预测

🍨 本文为🔗365天深度学习训练营中的学习记录博客
🍖 原作者：K同学啊

文章目录

[1. 简介](#1. 简介)
[2. 环境](#2. 环境)
[3. 代码实现](#3. 代码实现)
- [3.1 前期准备](#3.1 前期准备)
- - [3.1.1 导入库与数据](#3.1.1 导入库与数据)
  - [3.1.2 探索数值特征相关性](#3.1.2 探索数值特征相关性)
  - [3.1.3 探索分类特征与回归目标的关系](#3.1.3 探索分类特征与回归目标的关系)
  - [3.1.4 探索数值特征与回归目标的关系](#3.1.4 探索数值特征与回归目标的关系)
- [3.2 数据预处理](#3.2 数据预处理)
- - [3.2.1 处理缺失值](#3.2.1 处理缺失值)
  - [3.2.2 编码分类变量](#3.2.2 编码分类变量)
  - [3.2.3 划分训练集与测试集](#3.2.3 划分训练集与测试集)
  - [3.2.4 探索特征重要性](#3.2.4 探索特征重要性)
  - [3.2.5 标准化并转换为张量](#3.2.5 标准化并转换为张量)
  - [3.2.6 构建 DataLoader](#3.2.6 构建 DataLoader)
- [3.3 模型建立与训练](#3.3 模型建立与训练)
- - [3.3.1 设置计算设备](#3.3.1 设置计算设备)
  - [3.3.2 定义 LSTM 回归模型](#3.3.2 定义 LSTM 回归模型)
  - [3.3.3 定义训练函数](#3.3.3 定义训练函数)
  - [3.3.4 定义测试函数](#3.3.4 定义测试函数)
  - [3.3.5 配置优化器与学习率](#3.3.5 配置优化器与学习率)
  - [3.3.6 训练模型](#3.3.6 训练模型)
[4. 模型评估](#4. 模型评估)
- [4.1 可视化训练过程](#4.1 可视化训练过程)
- [4.2 最终指标](#4.2 最终指标)
[5. 总结](#5. 总结)

1. 简介

本实验使用 PyTorch 构建单层 LSTM 回归模型，根据年龄、BMI、吸烟状态、住院次数、保险类型、上一年度医疗费用等 19 个特征，预测个人年度医疗费用 annual_medical_cost。

项目	内容
模型	单层 LSTM（hidden_size=200）+ 全连接回归层
任务	年度医疗费用回归预测
数据集	5000 条记录、19 个输入特征，训练集/测试集按 80%/20% 划分
模型参数量	177,001，全部可训练
最终性能	测试集 R²=0.9808，MSE=934570.101，近似 RMSE=966.73

本实验将每条包含 19 个特征的记录转换为 (1, 19)，因此 LSTM 接收的序列长度为 1。它在这里主要充当带门控结构的非线性表格回归器，并未真正建模时间序列。

2. 环境

语言环境：Python 3.9.13（来自 Notebook 元数据）
编译器：Jupyter Notebook
深度学习框架：PyTorch
数据分析与可视化：Pandas、NumPy、Matplotlib、Seaborn
机器学习工具：scikit-learn
本次 Notebook 运行设备：CPU

3. 代码实现

3.1 前期准备

3.1.1 导入库与数据

首先导入数据处理、可视化、模型训练和评价所需的库，并读取医疗费用预测数据集。

python 复制代码

import warnings

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import torch
import torch.utils.data as data
from sklearn import metrics
from torch import nn

warnings.filterwarnings("ignore")

df = pd.read_csv("./Data/medical_cost_prediction_dataset.csv")
df.head(5)

前 5 条数据如下：

age	gender	bmi	smoker	diabetes	hypertension	insurance_type	insurance_coverage_pct	previous_year_cost	annual_medical_cost
69	Male	29.4	No	1	0	Private	80	10885	2645.50
32	Female	22.9	No	1	0	Government	64	18722	10959.70
89	Male	25.7	No	0	0	缺失	0	4196	8409.80
78	Male	31.9	Yes	0	1	Government	70	11128	7996.62
38	Male	27.7	No	0	0	Private	77	15110	3202.52

数据集共有 5000 行、20 列 ，其中前 19 列为特征，最后一列 annual_medical_cost 为回归目标。目标值均值为 8048.89，中位数为 5539.78，范围为 404.95～44792.10，整体呈现明显的右偏和长尾特征。

3.1.2 探索数值特征相关性

选择所有数值列并绘制 Pearson 相关系数热力图：

python 复制代码

numeric_cols = df.select_dtypes(include=["int64", "float64"])

plt.figure(figsize=(12, 8))
sns.heatmap(numeric_cols.corr(), cmap="coolwarm", annot=True)
plt.title("Correlation Heatmap")
plt.show()

从热力图可见：

insurance_coverage_pct 与年度医疗费用的线性相关性最强，相关系数约为 -0.86；
hospital_admissions 与目标呈中等程度正相关，相关系数约为 0.36；
medication_count、heart_disease 与目标的相关系数分别约为 0.13 和 0.12；
BMI、年龄、睡眠时间、每日步数等单个特征与目标的线性相关性较弱。

相关系数只能描述线性关系，不能直接代表因果关系；分类变量也不会出现在该数值热力图中，因此还需要结合分组图和模型特征重要性继续分析。

3.1.3 探索分类特征与回归目标的关系

使用箱线图观察年度医疗费用的总体分布，以及住院次数、吸烟状态与费用之间的关系。

python 复制代码

sns.set(font="SimHei", font_scale=0.8, style="darkgrid")
fig, ax = plt.subplots(1, 3, figsize=(12, 4))

sns.boxplot(data=df.loc[:, ["annual_medical_cost"]], ax=ax[0], whis=3)
ax[0].set_title("多个数值变量")

sns.boxplot(
    x=df["hospital_admissions"],
    y=df["annual_medical_cost"],
    ax=ax[1],
    whis=3,
)
ax[1].set_title("一个数值变量多个分组")

sns.boxplot(
    x="hospital_admissions",
    y="annual_medical_cost",
    hue="smoker",
    data=df,
    palette="Set1",
    width=0.5,
    ax=ax[2],
    whis=3,
)
ax[2].set_title("一个数值变量多个分组/子分组")

plt.tight_layout()
plt.show()

箱线图显示，年度医疗费用中存在较多高值样本；随着住院次数增加，费用中位数和分布上界整体上升。再使用小提琴图观察心脏病与吸烟状态分组下的分布形状：

python 复制代码

fig, ax = plt.subplots(1, 3, figsize=(12, 4))

sns.violinplot(data=df.loc[:, ["annual_medical_cost"]], ax=ax[0])
ax[0].set_title("多个数值变量")

sns.violinplot(
    x=df["heart_disease"],
    y=df["annual_medical_cost"],
    ax=ax[1],
)
ax[1].set_title("一个数值变量多个分组")

sns.violinplot(
    x="heart_disease",
    y="annual_medical_cost",
    hue="smoker",
    data=df,
    palette="Set1",
    width=0.5,
    ax=ax[2],
)
ax[2].set_title("一个数值变量多个分组/子分组")

plt.tight_layout()
plt.show()

3.1.4 探索数值特征与回归目标的关系

从前 100 条数据中选取上一年费用、BMI、睡眠时间等特征绘制气泡散点图，并用气泡大小表示年龄。

python 复制代码

fig, ax = plt.subplots(1, 3, figsize=(12, 4))

sns.scatterplot(
    data=df[:100],
    x="previous_year_cost",
    y="annual_medical_cost",
    size="age",
    sizes=(20, 200),
    alpha=0.6,
    color="#2E86AB",
    ax=ax[0],
)
ax[0].set_title("（气泡大小=年龄）")
ax[0].set_xlabel("上一年费用（元）")
ax[0].set_ylabel("年度医疗费用（元）")

sns.scatterplot(
    data=df[:100],
    x="bmi",
    y="annual_medical_cost",
    size="age",
    sizes=(40, 250),
    alpha=0.7,
    hue="gender",
    palette="Set2",
    legend=False,
    ax=ax[1],
)
ax[1].set_title("按 BMI 分组（气泡大小=年龄）")

sns.scatterplot(
    data=df[:100],
    x="sleep_hours",
    y="annual_medical_cost",
    hue="smoker",
    size="age",
    sizes=(40, 300),
    alpha=0.7,
    palette="Set1",
    ax=ax[2],
)
ax[2].set_title("睡眠时间+吸烟状态（气泡大小=年龄）")

plt.tight_layout()
plt.show()

散点图进一步印证了热力图中的结论：这些特征单独与年度医疗费用之间的关系较弱，模型需要综合多个特征捕捉非线性规律。

3.2 数据预处理

3.2.1 处理缺失值

python 复制代码

df.isnull().any()

检查结果表明，只有 insurance_type 存在缺失值，共 1048 条 。按照 Notebook 的处理方式，将其填充为字符串 "0"，用于表示缺失保险类别。

python 复制代码

df["insurance_type"] = df["insurance_type"].fillna("0")

X = df.iloc[:, :-1]
y = df.iloc[:, -1]

3.2.2 编码分类变量

数据中的 gender、smoker、physical_activity_level、insurance_type 和 city_type 为字符串类型。为了便于后续进行随机森林特征重要性排序，实验使用 OrdinalEncoder 将其转换为整数。

python 复制代码

from sklearn.preprocessing import OrdinalEncoder

label_cols = [1, 3, 8, 15, 17]
oe = OrdinalEncoder()
X.iloc[:, label_cols] = oe.fit_transform(X.iloc[:, label_cols])

整数编码会给无序类别引入人为的大小关系。若目标是获得更稳健的回归模型，可以在划分数据后只用训练集拟合编码器，并尝试 One-Hot 编码或类别嵌入。

3.2.3 划分训练集与测试集

python 复制代码

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.2,
    random_state=1,
)

X_train.shape, X_test.shape, y_train.shape, y_test.shape

text 复制代码

((4000, 19), (1000, 19), (4000,), (1000,))

3.2.4 探索特征重要性

使用随机森林回归器对训练集进行拟合，查看特征的重要性排行。

python 复制代码

from sklearn.ensemble import RandomForestRegressor

rf = RandomForestRegressor(n_estimators=100, random_state=42)
rf.fit(X_train, y_train)

feature_importance = pd.DataFrame({
    "feature": X_train.columns,
    "importance": rf.feature_importances_,
}).sort_values("importance", ascending=False)

print(feature_importance.head(10))

排名	特征	重要性
1	insurance_coverage_pct	0.428129
2	insurance_type	0.304765
3	hospital_admissions	0.189374
4	medication_count	0.018484
5	previous_year_cost	0.018046
6	heart_disease	0.015294
7	smoker	0.005553
8	doctor_visits_per_year	0.003000
9	diabetes	0.002681
10	bmi	0.002434

前三个特征的重要性之和约为 92.23%，说明保险覆盖比例、保险类型和住院次数对当前数据集的预测结果起主导作用。随机森林重要性并不等同于因果贡献，而且整数编码可能影响类别特征的重要性，因此该结果应作为探索性参考。

3.2.5 标准化并转换为张量

StandardScaler 将每个特征转换为均值约为 0、标准差约为 1 的标准分数，而不是将数值压缩到 0～1。

python 复制代码

import numpy as np
from sklearn.preprocessing import StandardScaler

sc = StandardScaler()
X_train = sc.fit_transform(X_train)
X_test = sc.transform(X_test)

X_train = torch.tensor(X_train, dtype=torch.float32).unsqueeze(1)
X_test = torch.tensor(X_test, dtype=torch.float32).unsqueeze(1)
y_train = torch.tensor(y_train.values, dtype=torch.float32)
y_test = torch.tensor(y_test.values, dtype=torch.float32)

X_train.shape, X_test.shape, y_train.shape, y_test.shape

text 复制代码

(torch.Size([4000, 1, 19]),
 torch.Size([1000, 1, 19]),
 torch.Size([4000]),
 torch.Size([1000]))

其中张量第二维的长度为 1，表示每条记录只有一个时间步；第三维的 19 表示输入特征数。

3.2.6 构建 DataLoader

python 复制代码

from torch.utils.data import DataLoader

batch_size = 32

train_dataset = data.TensorDataset(X_train, y_train)
test_dataset = data.TensorDataset(X_test, y_test)

train_dataloader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
)
test_dataloader = DataLoader(
    test_dataset,
    batch_size=batch_size,
)

3.3 模型建立与训练

3.3.1 设置计算设备

python 复制代码

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

text 复制代码

device(type='cpu')

3.3.2 定义 LSTM 回归模型

模型由一个输入维度为 19、隐藏层维度为 200 的单层 LSTM，以及一个将 200 维隐藏状态映射为单个回归值的全连接层构成。

组件	输入与输出	作用
LSTM	`(batch, 1, 19) → (batch, 1, 200)`	对 19 个输入特征进行门控非线性变换
Linear	`(batch, 1, 200) → (batch, 1, 1)`	输出年度医疗费用预测值

python 复制代码

class ModelLSTM(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=19,
            hidden_size=200,
            num_layers=1,
            batch_first=True,
        )
        self.fc = nn.Linear(200, 1)

    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.fc(out)
        return out


model = ModelLSTM().to(device)

Notebook 中通过 torchinfo.summary(model, (64, 1, 19)) 得到如下模型统计：

text 复制代码

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
ModelLSTM                               [64, 1, 1]                --
├─LSTM: 1-1                             [64, 1, 200]              176,800
├─Linear: 1-2                           [64, 1, 1]                201
==========================================================================================
Total params: 177,001
Trainable params: 177,001
Non-trainable params: 0
Total mult-adds (M): 11.33
Estimated Total Size (MB): 0.82
==========================================================================================

3.3.3 定义训练函数

训练阶段依次执行前向传播、MSE 损失计算、反向传播和参数更新，并在整个训练集上计算 R²。

python 复制代码

def train(dataloader, model, loss_fn, optimizer):
    num_batches = len(dataloader)
    train_loss = 0
    pred_list = []
    y_list = []

    for X, y in dataloader:
        X, y = X.to(device), y.to(device)

        pred = model(X).squeeze()
        loss = loss_fn(pred, y)

        y_list += [i.detach().numpy() for i in y.cpu()]
        pred_list += [i.detach().numpy() for i in pred.cpu()]

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    r2 = metrics.r2_score(y_list, pred_list)
    train_loss /= num_batches
    return r2, train_loss

3.3.4 定义测试函数

测试阶段使用 torch.no_grad() 关闭梯度计算，不更新模型参数，只统计测试集 MSE 和 R²。

python 复制代码

def test(dataloader, model, loss_fn):
    num_batches = len(dataloader)
    test_loss = 0
    pred_list = []
    y_list = []

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)

            pred = model(X).squeeze()
            loss = loss_fn(pred, y)

            y_list += [i.detach().numpy() for i in y.cpu()]
            pred_list += [i.detach().numpy() for i in pred.cpu()]
            test_loss += loss.item()

    r2 = metrics.r2_score(y_list, pred_list)
    test_loss /= num_batches
    return r2, test_loss

3.3.5 配置优化器与学习率

本实验使用 Adam 优化器和均方误差损失。初始学习率为 0.1，每经过 5 个 Epoch，学习率衰减为原来的 0.92。

python 复制代码

def adjust_learning_rate(optimizer, epoch, start_lr):
    lr = start_lr * (0.92 ** (epoch // 5))
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr


learn_rate = 0.1
optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate)
loss_fn = nn.MSELoss()
epochs = 50

3.3.6 训练模型

python 复制代码

train_loss = []
train_r2 = []
test_loss = []
test_r2 = []

for epoch in range(epochs):
    adjust_learning_rate(optimizer, epoch, learn_rate)

    model.train()
    epoch_train_r2, epoch_train_loss = train(
        train_dataloader,
        model,
        loss_fn,
        optimizer,
    )

    model.eval()
    epoch_test_r2, epoch_test_loss = test(
        test_dataloader,
        model,
        loss_fn,
    )

    train_r2.append(epoch_train_r2)
    train_loss.append(epoch_train_loss)
    test_r2.append(epoch_test_r2)
    test_loss.append(epoch_test_loss)

    lr = optimizer.param_groups[0]["lr"]
    print(
        f"Epoch:{epoch + 1:2d}, "
        f"Train_R2:{epoch_train_r2:.3f}, "
        f"Train_loss:{epoch_train_loss:.3f}, "
        f"Test_R2:{epoch_test_r2:.4f}, "
        f"Test_loss:{epoch_test_loss:.3f}, "
        f"Lr:{lr:.2E}"
    )

print("Done")

50 轮训练中的关键节点如下：

Epoch	Train R²	Train MSE	Test R²	Test MSE	学习率
1	-1.005	100338544.272	-0.7290	85010290.875	1.00E-01
5	0.139	43067673.288	0.2080	38662337.234	1.00E-01
10	0.581	20965871.826	0.6066	19125876.430	9.20E-02
15	0.772	11431778.461	0.7823	10586175.012	8.46E-02
20	0.866	6705683.027	0.8697	6336887.557	7.79E-02
25	0.916	4191079.665	0.9155	4109939.980	7.16E-02
30	0.945	2754931.301	0.9420	2823772.773	6.59E-02
35	0.962	1886100.701	0.9583	2028954.141	6.06E-02
40	0.973	1338526.250	0.9692	1503044.938	5.58E-02
45	0.980	983236.457	0.9764	1154301.055	5.13E-02
50	0.985	742104.021	0.9808	934570.101	4.72E-02

4. 模型评估

4.1 可视化训练过程

下面绘制训练集与测试集的 R²、MSE 随 Epoch 的变化。由于 MSE 从约 10^8 下降至 10^6 以下，损失图使用对数纵轴以便观察完整收敛过程。

python 复制代码

epochs_range = range(1, epochs + 1)

plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_r2, label="Train R²")
plt.plot(epochs_range, test_r2, label="Test R²")
plt.axhline(0, color="gray", linewidth=0.8, linestyle="--")
plt.xlabel("Epoch")
plt.ylabel("R²")
plt.title("Training and Test R²")
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss, label="Train MSE")
plt.plot(epochs_range, test_loss, label="Test MSE")
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("MSE (log scale)")
plt.title("Training and Test Loss")
plt.legend()

plt.tight_layout()
plt.show()

从曲线可以看到：

前 3 个 Epoch 的 R² 为负，说明初始模型还不如直接使用目标均值进行预测；测试集 R² 在第 4 个 Epoch 转为正值。
模型随后稳定收敛，测试集 R² 在第 23 个 Epoch 首次超过 0.90，并在第 50 个 Epoch 达到 0.9808。
训练集和测试集曲线非常接近，最终 R² 差值约为 0.0042，当前单次划分下未出现明显过拟合。
训练与测试 MSE 均持续下降，说明当前 50 轮训练尚未出现性能反弹。

4.2 最终指标

指标	训练集	测试集
R²	0.985	0.9808
MSE	742104.021	934570.101
由批平均 MSE 估算的 RMSE	约 861.45	约 966.73

测试集 R² 为 0.9808 ，表示模型在该测试划分上能够解释约 98.08% 的目标方差。由于 Notebook 中的 test_loss 是各批次 MSE 的简单平均，而最后一个批次样本量略小，表中的 RMSE 是对该记录值开平方得到的近似结果，并非重新按全部样本计算的精确全局 RMSE。

5. 总结

模型拟合效果较好 ：经过 50 个 Epoch，测试集 R² 达到 0.9808 ，由测试 MSE 估算的 RMSE 约为 966.73。相较于目标均值 8048.89，模型在当前数据划分上取得了较好的预测效果。
少数特征贡献突出 ：随机森林结果显示，insurance_coverage_pct、insurance_type 和 hospital_admissions 的重要性合计约 92.23%。这也解释了模型为何能够较快获得较高的 R²，但这种集中度需要在其他数据或交叉验证中进一步确认。
模型参数量适中 ：该 LSTM 共有 177,001 个可训练参数，训练集和测试集曲线接近，当前实验中未观察到明显过拟合。
LSTM 的结构优势尚未充分发挥：输入序列长度只有 1，19 个字段被作为同一时间步的特征输入，因此不存在跨时间步依赖。对于这种表格数据，建议增加线性回归、随机森林、梯度提升树和多层感知机作为基线，比较精度、训练成本与可解释性。
评估流程仍可改进：当前每轮都查看测试集指标，实际上测试集承担了验证集角色。后续可划分独立的训练集、验证集和测试集，使用验证集选取最佳模型，只在最后一次评估时使用测试集。
提升可复现性与稳健性：建议固定 NumPy、PyTorch 和 DataLoader 的随机种子；只在训练集上拟合编码器与标准化器；使用 One-Hot 编码或类别嵌入；并通过 K 折交叉验证检验结果是否稳定。