Skip to content

Commit 0dc0b5f

Browse files
EricSchrockEricSchrock
authored
Update MIMIC3 example python scripts (#834)
1 parent e01a659 commit 0dc0b5f

10 files changed

Lines changed: 291 additions & 295 deletions

‎examples/drug_recommendation/drug_recommendation_mimic3_gamenet.py‎

Lines changed: 44 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,49 @@
1+
import tempfile
2+
13
from pyhealth.datasets import MIMIC3Dataset
24
from pyhealth.datasets import split_by_patient, get_dataloader
35
from pyhealth.models import GAMENet
4-
from pyhealth.tasks import drug_recommendation_mimic3_fn
6+
from pyhealth.tasks import DrugRecommendationMIMIC3
57
from pyhealth.trainer import Trainer
68

7-
# STEP 1: load data
8-
base_dataset = MIMIC3Dataset(
9-
root="/srv/local/data/physionet.org/files/mimiciii/1.4",
10-
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
11-
code_mapping={"NDC": ("ATC", {"target_kwargs": {"level": 3}})},
12-
dev=True,
13-
refresh_cache=False,
14-
)
15-
base_dataset.stat()
16-
17-
# STEP 2: set task
18-
sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
19-
sample_dataset.stat()
20-
21-
train_dataset, val_dataset, test_dataset = split_by_patient(
22-
sample_dataset, [0.8, 0.1, 0.1]
23-
)
24-
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
25-
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
26-
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
27-
28-
# STEP 3: define model
29-
model = GAMENet(
30-
sample_dataset,
31-
)
32-
33-
# STEP 4: define trainer
34-
trainer = Trainer(
35-
model=model,
36-
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples", "ddi"],
37-
)
38-
39-
trainer.train(
40-
train_dataloader=train_dataloader,
41-
val_dataloader=val_dataloader,
42-
epochs=20,
43-
monitor="pr_auc_samples",
44-
)
45-
46-
# STEP 5: evaluate
47-
print (trainer.evaluate(test_dataloader))
9+
if __name__ == "__main__":
10+
# STEP 1: load data
11+
base_dataset = MIMIC3Dataset(
12+
root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III",
13+
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
14+
cache_dir=tempfile.TemporaryDirectory().name,
15+
dev=True,
16+
)
17+
base_dataset.stats()
18+
19+
# STEP 2: set task
20+
task = DrugRecommendationMIMIC3()
21+
sample_dataset = base_dataset.set_task(task)
22+
23+
train_dataset, val_dataset, test_dataset = split_by_patient(
24+
sample_dataset, [0.8, 0.1, 0.1]
25+
)
26+
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
27+
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
28+
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
29+
30+
# STEP 3: define model
31+
model = GAMENet(
32+
sample_dataset,
33+
)
34+
35+
# STEP 4: define trainer
36+
trainer = Trainer(
37+
model=model,
38+
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples", "ddi"],
39+
)
40+
41+
trainer.train(
42+
train_dataloader=train_dataloader,
43+
val_dataloader=val_dataloader,
44+
epochs=1,
45+
monitor="pr_auc_samples",
46+
)
47+
48+
# STEP 5: evaluate
49+
print(trainer.evaluate(test_dataloader))

‎examples/drug_recommendation/drug_recommendation_mimic3_molerec.py‎

Lines changed: 44 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,50 +1,49 @@
1+
import tempfile
2+
13
from pyhealth.datasets import MIMIC3Dataset
24
from pyhealth.datasets import split_by_patient, get_dataloader
35
from pyhealth.models import MoleRec
4-
from pyhealth.tasks import drug_recommendation_mimic3_fn
6+
from pyhealth.tasks import DrugRecommendationMIMIC3
57
from pyhealth.trainer import Trainer
68

7-
# STEP 1: load data
8-
base_dataset = MIMIC3Dataset(
9-
root="/srv/local/data/physionet.org/files/mimiciii/1.4",
10-
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
11-
code_mapping={"NDC": ("ATC", {"target_kwargs": {"level": 3}})},
12-
dev=True,
13-
refresh_cache=False,
14-
)
15-
base_dataset.stat()
16-
17-
# STEP 2: set task
18-
sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
19-
sample_dataset.stat()
20-
21-
train_dataset, val_dataset, test_dataset = split_by_patient(
22-
sample_dataset, [0.8, 0.1, 0.1]
23-
)
24-
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
25-
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
26-
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
27-
28-
# STEP 3: define model
29-
model = MoleRec(
30-
sample_dataset,
31-
feature_keys=["conditions", "procedures"],
32-
label_key="drugs",
33-
mode="multilabel",
34-
)
35-
36-
# STEP 4: define trainer
37-
trainer = Trainer(
38-
model=model,
39-
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples", "ddi"],
40-
)
41-
42-
trainer.train(
43-
train_dataloader=train_dataloader,
44-
val_dataloader=val_dataloader,
45-
epochs=3,
46-
monitor="pr_auc_samples",
47-
)
48-
49-
# STEP 5: evaluate
50-
print (trainer.evaluate(test_dataloader))
9+
if __name__ == "__main__":
10+
# STEP 1: load data
11+
base_dataset = MIMIC3Dataset(
12+
root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III",
13+
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
14+
cache_dir=tempfile.TemporaryDirectory().name,
15+
dev=True,
16+
)
17+
base_dataset.stats()
18+
19+
# STEP 2: set task
20+
task = DrugRecommendationMIMIC3()
21+
sample_dataset = base_dataset.set_task(task)
22+
23+
train_dataset, val_dataset, test_dataset = split_by_patient(
24+
sample_dataset, [0.8, 0.1, 0.1]
25+
)
26+
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
27+
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
28+
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
29+
30+
# STEP 3: define model
31+
model = MoleRec(
32+
sample_dataset,
33+
)
34+
35+
# STEP 4: define trainer
36+
trainer = Trainer(
37+
model=model,
38+
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples", "ddi"],
39+
)
40+
41+
trainer.train(
42+
train_dataloader=train_dataloader,
43+
val_dataloader=val_dataloader,
44+
epochs=1,
45+
monitor="pr_auc_samples",
46+
)
47+
48+
# STEP 5: evaluate
49+
print(trainer.evaluate(test_dataloader))

‎examples/drug_recommendation/drug_recommendation_mimic3_safedrug.py‎

Lines changed: 44 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,49 @@
1+
import tempfile
2+
13
from pyhealth.datasets import MIMIC3Dataset
24
from pyhealth.datasets import split_by_patient, get_dataloader
35
from pyhealth.models import SafeDrug
4-
from pyhealth.tasks import drug_recommendation_mimic3_fn
6+
from pyhealth.tasks import DrugRecommendationMIMIC3
57
from pyhealth.trainer import Trainer
68

7-
# STEP 1: load data
8-
base_dataset = MIMIC3Dataset(
9-
root="/srv/local/data/physionet.org/files/mimiciii/1.4",
10-
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
11-
code_mapping={"NDC": ("ATC", {"target_kwargs": {"level": 3}})},
12-
dev=True,
13-
refresh_cache=False,
14-
)
15-
base_dataset.stat()
16-
17-
# STEP 2: set task
18-
sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
19-
sample_dataset.stat()
20-
21-
train_dataset, val_dataset, test_dataset = split_by_patient(
22-
sample_dataset, [0.8, 0.1, 0.1]
23-
)
24-
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
25-
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
26-
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
27-
28-
# STEP 3: define model
29-
model = SafeDrug(
30-
sample_dataset,
31-
)
32-
33-
# STEP 4: define trainer
34-
trainer = Trainer(
35-
model=model,
36-
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples", "ddi"],
37-
)
38-
39-
trainer.train(
40-
train_dataloader=train_dataloader,
41-
val_dataloader=val_dataloader,
42-
epochs=25,
43-
monitor="pr_auc_samples",
44-
)
45-
46-
# STEP 5: evaluate
47-
print (trainer.evaluate(test_dataloader))
9+
if __name__ == "__main__":
10+
# STEP 1: load data
11+
base_dataset = MIMIC3Dataset(
12+
root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III",
13+
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
14+
cache_dir=tempfile.TemporaryDirectory().name,
15+
dev=True,
16+
)
17+
base_dataset.stats()
18+
19+
# STEP 2: set task
20+
task = DrugRecommendationMIMIC3()
21+
sample_dataset = base_dataset.set_task(task)
22+
23+
train_dataset, val_dataset, test_dataset = split_by_patient(
24+
sample_dataset, [0.8, 0.1, 0.1]
25+
)
26+
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
27+
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
28+
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
29+
30+
# STEP 3: define model
31+
model = SafeDrug(
32+
sample_dataset,
33+
)
34+
35+
# STEP 4: define trainer
36+
trainer = Trainer(
37+
model=model,
38+
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples", "ddi"],
39+
)
40+
41+
trainer.train(
42+
train_dataloader=train_dataloader,
43+
val_dataloader=val_dataloader,
44+
epochs=1,
45+
monitor="pr_auc_samples",
46+
)
47+
48+
# STEP 5: evaluate
49+
print(trainer.evaluate(test_dataloader))

‎examples/drug_recommendation/drug_recommendation_mimic3_transformer.py‎

Lines changed: 44 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,49 @@
1+
import tempfile
2+
13
from pyhealth.datasets import MIMIC3Dataset
24
from pyhealth.datasets import split_by_patient, get_dataloader
35
from pyhealth.models import Transformer
4-
from pyhealth.tasks import drug_recommendation_mimic3_fn
6+
from pyhealth.tasks import DrugRecommendationMIMIC3
57
from pyhealth.trainer import Trainer
68

7-
# STEP 1: load data
8-
base_dataset = MIMIC3Dataset(
9-
root="/srv/local/data/physionet.org/files/mimiciii/1.4",
10-
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
11-
code_mapping={"NDC": ("ATC", {"target_kwargs": {"level": 3}})},
12-
dev=True,
13-
refresh_cache=True,
14-
)
15-
16-
base_dataset.stat()
17-
18-
# STEP 2: set task
19-
sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
20-
sample_dataset.stat()
21-
22-
train_dataset, val_dataset, test_dataset = split_by_patient(
23-
sample_dataset, [0.8, 0.1, 0.1]
24-
)
25-
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
26-
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
27-
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
28-
29-
# STEP 3: define model
30-
model = Transformer(
31-
dataset=sample_dataset,
32-
feature_keys=["conditions", "procedures"],
33-
label_key="drugs",
34-
mode="multilabel",
35-
)
36-
37-
# STEP 4: define trainer
38-
trainer = Trainer(
39-
model=model,
40-
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples"],
41-
)
42-
43-
trainer.train(
44-
train_dataloader=train_dataloader,
45-
val_dataloader=val_dataloader,
46-
epochs=20,
47-
monitor="pr_auc_samples",
48-
)
49-
50-
# STEP 5: evaluate
51-
print (trainer.evaluate(test_dataloader))
9+
if __name__ == "__main__":
10+
# STEP 1: load data
11+
base_dataset = MIMIC3Dataset(
12+
root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III",
13+
tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
14+
cache_dir=tempfile.TemporaryDirectory().name,
15+
dev=True,
16+
)
17+
base_dataset.stats()
18+
19+
# STEP 2: set task
20+
task = DrugRecommendationMIMIC3()
21+
sample_dataset = base_dataset.set_task(task)
22+
23+
train_dataset, val_dataset, test_dataset = split_by_patient(
24+
sample_dataset, [0.8, 0.1, 0.1]
25+
)
26+
train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
27+
val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
28+
test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
29+
30+
# STEP 3: define model
31+
model = Transformer(
32+
dataset=sample_dataset,
33+
)
34+
35+
# STEP 4: define trainer
36+
trainer = Trainer(
37+
model=model,
38+
metrics=["jaccard_samples", "f1_samples", "pr_auc_samples"],
39+
)
40+
41+
trainer.train(
42+
train_dataloader=train_dataloader,
43+
val_dataloader=val_dataloader,
44+
epochs=1,
45+
monitor="pr_auc_samples",
46+
)
47+
48+
# STEP 5: evaluate
49+
print(trainer.evaluate(test_dataloader))

0 commit comments

Comments
 (0)