"""Functions for training and running EF prediction."""
import math
import os
import time
import click
import matplotlib.pyplot as plt
import numpy as np
import sklearn.metrics
import torch
import torchvision
import tqdm
import echonet
@click.command("video")
@click.option("--data_dir", type=click.Path(exists=True, file_okay=False), default=None)
@click.option("--output", type=click.Path(file_okay=False), default=None)
@click.option("--task", type=str, default="EF")
@click.option("--model_name", type=click.Choice(
sorted(name for name in torchvision.models.video.__dict__
if name.islower() and not name.startswith("__") and callable(torchvision.models.video.__dict__[name]))),
default="r2plus1d_18")
@click.option("--pretrained/--random", default=True)
@click.option("--weights", type=click.Path(exists=True, dir_okay=False), default=None)
@click.option("--run_test/--skip_test", default=False)
@click.option("--num_epochs", type=int, default=45)
@click.option("--lr", type=float, default=1e-4)
@click.option("--weight_decay", type=float, default=1e-4)
@click.option("--lr_step_period", type=int, default=15)
@click.option("--frames", type=int, default=32)
@click.option("--period", type=int, default=2)
@click.option("--num_train_patients", type=int, default=None)
@click.option("--num_workers", type=int, default=4)
@click.option("--batch_size", type=int, default=20)
@click.option("--device", type=str, default=None)
@click.option("--seed", type=int, default=0)
def run(
data_dir=None,
output=None,
task="EF",
model_name="r2plus1d_18",
pretrained=True,
weights=None,
run_test=False,
num_epochs=45,
lr=1e-4,
weight_decay=1e-4,
lr_step_period=15,
frames=32,
period=2,
num_train_patients=None,
num_workers=4,
batch_size=20,
device=None,
seed=0,
):
"""Trains/tests EF prediction model.
\b
Args:
data_dir (str, optional): Directory containing dataset. Defaults to
`echonet.config.DATA_DIR`.
output (str, optional): Directory to place outputs. Defaults to
output/video/<model_name>_<pretrained/random>/.
task (str, optional): Name of task to predict. Options are the headers
of FileList.csv. Defaults to ``EF''.
model_name (str, optional): Name of model. One of ``mc3_18'',
``r2plus1d_18'', or ``r3d_18''
(options are torchvision.models.video.<model_name>)
Defaults to ``r2plus1d_18''.
pretrained (bool, optional): Whether to use pretrained weights for model
Defaults to True.
weights (str, optional): Path to checkpoint containing weights to
initialize model. Defaults to None.
run_test (bool, optional): Whether or not to run on test.
Defaults to False.
num_epochs (int, optional): Number of epochs during training.
Defaults to 45.
lr (float, optional): Learning rate for SGD
Defaults to 1e-4.
weight_decay (float, optional): Weight decay for SGD
Defaults to 1e-4.
lr_step_period (int or None, optional): Period of learning rate decay
(learning rate is decayed by a multiplicative factor of 0.1)
Defaults to 15.
frames (int, optional): Number of frames to use in clip
Defaults to 32.
period (int, optional): Sampling period for frames
Defaults to 2.
n_train_patients (int or None, optional): Number of training patients
for ablations. Defaults to all patients.
num_workers (int, optional): Number of subprocesses to use for data
loading. If 0, the data will be loaded in the main process.
Defaults to 4.
device (str or None, optional): Name of device to run on. Options from
https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device
Defaults to ``cuda'' if available, and ``cpu'' otherwise.
batch_size (int, optional): Number of samples to load per batch
Defaults to 20.
seed (int, optional): Seed for random number generator. Defaults to 0.
"""
# Seed RNGs
np.random.seed(seed)
torch.manual_seed(seed)
# Set default output directory
if output is None:
output = os.path.join("output", "video", "{}_{}_{}_{}".format(model_name, frames, period, "pretrained" if pretrained else "random"))
os.makedirs(output, exist_ok=True)
# Set device for computations
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Set up model
model = torchvision.models.video.__dict__[model_name](pretrained=pretrained)
model.fc = torch.nn.Linear(model.fc.in_features, 1)
model.fc.bias.data[0] = 55.6
if device.type == "cuda":
model = torch.nn.DataParallel(model)
model.to(device)
if weights is not None:
checkpoint = torch.load(weights)
model.load_state_dict(checkpoint['state_dict'])
# Set up optimizer
optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay)
if lr_step_period is None:
lr_step_period = math.inf
scheduler = torch.optim.lr_scheduler.StepLR(optim, lr_step_period)
# Compute mean and std
mean, std = echonet.utils.get_mean_and_std(echonet.datasets.Echo(root=data_dir, split="train"))
kwargs = {"target_type": task,
"mean": mean,
"std": std,
"length": frames,
"period": period,
}
# Set up datasets and dataloaders
dataset = {}
dataset["train"] = echonet.datasets.Echo(root=data_dir, split="train", **kwargs, pad=12)
if num_train_patients is not None and len(dataset["train"]) > num_train_patients:
# Subsample patients (used for ablation experiment)
indices = np.random.choice(len(dataset["train"]), num_train_patients, replace=False)
dataset["train"] = torch.utils.data.Subset(dataset["train"], indices)
dataset["val"] = echonet.datasets.Echo(root=data_dir, split="val", **kwargs)
# Run training and testing loops
with open(os.path.join(output, "log.csv"), "a") as f:
epoch_resume = 0
bestLoss = float("inf")
try:
# Attempt to load checkpoint
checkpoint = torch.load(os.path.join(output, "checkpoint.pt"))
model.load_state_dict(checkpoint['state_dict'])
optim.load_state_dict(checkpoint['opt_dict'])
scheduler.load_state_dict(checkpoint['scheduler_dict'])
epoch_resume = checkpoint["epoch"] + 1
bestLoss = checkpoint["best_loss"]
f.write("Resuming from epoch {}\n".format(epoch_resume))
except FileNotFoundError:
f.write("Starting run from scratch\n")
for epoch in range(epoch_resume, num_epochs):
print("Epoch #{}".format(epoch), flush=True)
for phase in ['train', 'val']:
start_time = time.time()
for i in range(torch.cuda.device_count()):
torch.cuda.reset_peak_memory_stats(i)
ds = dataset[phase]
dataloader = torch.utils.data.DataLoader(
ds, batch_size=batch_size, num_workers=num_workers, shuffle=True, pin_memory=(device.type == "cuda"), drop_last=(phase == "train"))
loss, yhat, y = echonet.utils.video.run_epoch(model, dataloader, phase == "train", optim, device)
f.write("{},{},{},{},{},{},{},{},{}\n".format(epoch,
phase,
loss,
sklearn.metrics.r2_score(y, yhat),
time.time() - start_time,
y.size,
sum(torch.cuda.max_memory_allocated() for i in range(torch.cuda.device_count())),
sum(torch.cuda.max_memory_reserved() for i in range(torch.cuda.device_count())),
batch_size))
f.flush()
scheduler.step()
# Save checkpoint
save = {
'epoch': epoch,
'state_dict': model.state_dict(),
'period': period,
'frames': frames,
'best_loss': bestLoss,
'loss': loss,
'r2': sklearn.metrics.r2_score(y, yhat),
'opt_dict': optim.state_dict(),
'scheduler_dict': scheduler.state_dict(),
}
torch.save(save, os.path.join(output, "checkpoint.pt"))
if loss < bestLoss:
torch.save(save, os.path.join(output, "best.pt"))
bestLoss = loss
# Load best weights
if num_epochs != 0:
checkpoint = torch.load(os.path.join(output, "best.pt"))
model.load_state_dict(checkpoint['state_dict'])
f.write("Best validation loss {} from epoch {}\n".format(checkpoint["loss"], checkpoint["epoch"]))
f.flush()
if run_test:
for split in ["val", "test"]:
# Performance without test-time augmentation
dataloader = torch.utils.data.DataLoader(
echonet.datasets.Echo(root=data_dir, split=split, **kwargs),
batch_size=batch_size, num_workers=num_workers, shuffle=True, pin_memory=(device.type == "cuda"))
loss, yhat, y = echonet.utils.video.run_epoch(model, dataloader, False, None, device)
f.write("{} (one clip) R2: {:.3f} ({:.3f} - {:.3f})\n".format(split, *echonet.utils.bootstrap(y, yhat, sklearn.metrics.r2_score)))
f.write("{} (one clip) MAE: {:.2f} ({:.2f} - {:.2f})\n".format(split, *echonet.utils.bootstrap(y, yhat, sklearn.metrics.mean_absolute_error)))
f.write("{} (one clip) RMSE: {:.2f} ({:.2f} - {:.2f})\n".format(split, *tuple(map(math.sqrt, echonet.utils.bootstrap(y, yhat, sklearn.metrics.mean_squared_error)))))
f.flush()
# Performance with test-time augmentation
ds = echonet.datasets.Echo(root=data_dir, split=split, **kwargs, clips="all")
dataloader = torch.utils.data.DataLoader(
ds, batch_size=1, num_workers=num_workers, shuffle=False, pin_memory=(device.type == "cuda"))
loss, yhat, y = echonet.utils.video.run_epoch(model, dataloader, False, None, device, save_all=True, block_size=batch_size)
f.write("{} (all clips) R2: {:.3f} ({:.3f} - {:.3f})\n".format(split, *echonet.utils.bootstrap(y, np.array(list(map(lambda x: x.mean(), yhat))), sklearn.metrics.r2_score)))
f.write("{} (all clips) MAE: {:.2f} ({:.2f} - {:.2f})\n".format(split, *echonet.utils.bootstrap(y, np.array(list(map(lambda x: x.mean(), yhat))), sklearn.metrics.mean_absolute_error)))
f.write("{} (all clips) RMSE: {:.2f} ({:.2f} - {:.2f})\n".format(split, *tuple(map(math.sqrt, echonet.utils.bootstrap(y, np.array(list(map(lambda x: x.mean(), yhat))), sklearn.metrics.mean_squared_error)))))
f.flush()
# Write full performance to file
with open(os.path.join(output, "{}_predictions.csv".format(split)), "w") as g:
for (filename, pred) in zip(ds.fnames, yhat):
for (i, p) in enumerate(pred):
g.write("{},{},{:.4f}\n".format(filename, i, p))
echonet.utils.latexify()
yhat = np.array(list(map(lambda x: x.mean(), yhat)))
# Plot actual and predicted EF
fig = plt.figure(figsize=(3, 3))
lower = min(y.min(), yhat.min())
upper = max(y.max(), yhat.max())
plt.scatter(y, yhat, color="k", s=1, edgecolor=None, zorder=2)
plt.plot([0, 100], [0, 100], linewidth=1, zorder=3)
plt.axis([lower - 3, upper + 3, lower - 3, upper + 3])
plt.gca().set_aspect("equal", "box")
plt.xlabel("Actual EF (%)")
plt.ylabel("Predicted EF (%)")
plt.xticks([10, 20, 30, 40, 50, 60, 70, 80])
plt.yticks([10, 20, 30, 40, 50, 60, 70, 80])
plt.grid(color="gainsboro", linestyle="--", linewidth=1, zorder=1)
plt.tight_layout()
plt.savefig(os.path.join(output, "{}_scatter.pdf".format(split)))
plt.close(fig)
# Plot AUROC
fig = plt.figure(figsize=(3, 3))
plt.plot([0, 1], [0, 1], linewidth=1, color="k", linestyle="--")
for thresh in [35, 40, 45, 50]:
fpr, tpr, _ = sklearn.metrics.roc_curve(y > thresh, yhat)
print(thresh, sklearn.metrics.roc_auc_score(y > thresh, yhat))
plt.plot(fpr, tpr)
plt.axis([-0.01, 1.01, -0.01, 1.01])
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.tight_layout()
plt.savefig(os.path.join(output, "{}_roc.pdf".format(split)))
plt.close(fig)
def run_epoch(model, dataloader, train, optim, device, save_all=False, block_size=None):
"""Run one epoch of training/evaluation for segmentation.
Args:
model (torch.nn.Module): Model to train/evaulate.
dataloder (torch.utils.data.DataLoader): Dataloader for dataset.
train (bool): Whether or not to train model.
optim (torch.optim.Optimizer): Optimizer
device (torch.device): Device to run on
save_all (bool, optional): If True, return predictions for all
test-time augmentations separately. If False, return only
the mean prediction.
Defaults to False.
block_size (int or None, optional): Maximum number of augmentations
to run on at the same time. Use to limit the amount of memory
used. If None, always run on all augmentations simultaneously.
Default is None.
"""
model.train(train)
total = 0 # total training loss
n = 0 # number of videos processed
s1 = 0 # sum of ground truth EF
s2 = 0 # Sum of ground truth EF squared
yhat = []
y = []
with torch.set_grad_enabled(train):
with tqdm.tqdm(total=len(dataloader)) as pbar:
for (X, outcome) in dataloader:
y.append(outcome.numpy())
X = X.to(device)
outcome = outcome.to(device)
average = (len(X.shape) == 6)
if average:
batch, n_clips, c, f, h, w = X.shape
X = X.view(-1, c, f, h, w)
s1 += outcome.sum()
s2 += (outcome ** 2).sum()
if block_size is None:
outputs = model(X)
else:
outputs = torch.cat([model(X[j:(j + block_size), ...]) for j in range(0, X.shape[0], block_size)])
if save_all:
yhat.append(outputs.view(-1).to("cpu").detach().numpy())
if average:
outputs = outputs.view(batch, n_clips, -1).mean(1)
if not save_all:
yhat.append(outputs.view(-1).to("cpu").detach().numpy())
loss = torch.nn.functional.mse_loss(outputs.view(-1), outcome)
if train:
optim.zero_grad()
loss.backward()
optim.step()
total += loss.item() * X.size(0)
n += X.size(0)
pbar.set_postfix_str("{:.2f} ({:.2f}) / {:.2f}".format(total / n, loss.item(), s2 / n - (s1 / n) ** 2))
pbar.update()
if not save_all:
yhat = np.concatenate(yhat)
y = np.concatenate(y)
return total / n, yhat, y
Datasets
Models
