def get_temporal_dataset(
dataset_file_path,
file_name_list,
batch_size,
balance,
shuffle,
workers_num,
collate_fn,
):
dataset = TemporalDataset(
dataset_file_path=dataset_file_path, file_name_list=file_name_list
)
if balance:
dataloader = DataLoader(
dataset=dataset,
batch_size=batch_size,
sampler=sampler.ImbalancedDatasetSampler(dataset),
num_workers=workers_num,
collate_fn=collate_fn,
pin_memory=False,
)
else:
dataloader = DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=shuffle,
num_workers=workers_num,
collate_fn=collate_fn,
pin_memory=False,
)
return dataloader
def mlp_cls(args):
input_size = 269
hidden_size = args.hidden_size
num_epochs = args.max_epoch
batch_size = args.batch_size
learning_rate = args.lr
w_decay = args.weight_decay
log_dir = args.save_result_root
imbalanced = args.sampler
target_type = args.target_type
save_result = True
sbp_target_idx = -2
dbp_target_idx = -1
sbp_num_class = 6
dbp_num_class = 5
output_size = dbp_num_class + sbp_num_class
stats = {
'sbp_mean': 132.28392012972589,
'dbp_mean': 72.38757001151521,
'sbp_std': 26.86375195359048,
'dbp_std': 14.179178540137421
}
writer = SummaryWriter(log_dir + 'logs/')
# device = torch.device('cpu')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MLP(input_size, hidden_size, output_size).to(device)
train_data = torch.load('data/MLP/Train.pt')
X = train_data[:, :-4]
y = train_data[:, [sbp_target_idx,
dbp_target_idx]] #Shape : (batch,2) --> (batch,1,1)
val_data = torch.load('data/MLP/Validation.pt')
X_val = val_data[:, :-4]
y_val = val_data[:, [sbp_target_idx, dbp_target_idx]]
train_dataset = loader.HD_Dataset((X, y))
imbalanced_sampler = sampler.ImbalancedDatasetSampler(y, target_type)
val_dataset = loader.HD_Dataset((X_val, y_val))
if imbalanced:
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
sampler=imbalanced_sampler)
else:
train_loader = DataLoader(train_dataset, batch_size=batch_size)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
sbp_criterion = nn.CrossEntropyLoss()
dbp_criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),
lr=learning_rate,
weight_decay=w_decay)
print("Starting training...")
total_step = len(train_loader)
best_acc = 0
for epoch in range(num_epochs):
sbp_correct = 0
dbp_correct = 0
total = 0
for i, (inputs, targets) in enumerate(train_loader):
# print("Input shape", inputs.shape)
# print("Targets shape", targets.shape)
inputs = inputs.float().to(device)
targets = targets.long().to(device)
# Forward pass
outputs = model(inputs)
sbp_outputs = outputs[:, :sbp_num_class]
dbp_outputs = outputs[:, sbp_num_class:]
_, sbp_pred = torch.max(sbp_outputs, 1)
_, dbp_pred = torch.max(dbp_outputs, 1)
sbp_loss = sbp_criterion(sbp_outputs, targets[:, 0])
dbp_loss = dbp_criterion(dbp_outputs, targets[:, 1])
total_loss = sbp_loss + dbp_loss
total += inputs.size(0)
sbp_correct += (targets[:, 0] == sbp_pred).sum().item()
dbp_correct += (targets[:, 1] == dbp_pred).sum().item()
# Backward and optimize
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if (i + 1) % args.valid_iter_freq == 0:
iteration = (i + 1) + (total_step * epoch)
print(
'\n\nEpoch [{}/{}], Step [{}/{}], SBP_Acc: {:.4f}, DBP_Acc: {:.4f}'
.format(epoch + 1, num_epochs, i + 1, total_step,
sbp_correct / total, dbp_correct / total))
writer.add_scalars('Loss/Train', {
'Total': total_loss,
'SBP': sbp_loss,
'DBP': dbp_loss
}, iteration)
writer.add_scalars('Acc/Train', {
'SBP': sbp_correct / total,
'DBP': dbp_correct / total
}, iteration)
with torch.no_grad():
model.eval()
val_total = 0
val_sbp_correct = 0
val_dbp_correct = 0
sbp_pred_tensor = torch.Tensor().long().to(device)
dbp_pred_tensor = torch.Tensor().long().to(device)
target_tensor = torch.Tensor().long().to(device)
for (inputs, targets) in val_loader:
inputs = inputs.float().to(device)
targets = targets.long().to(device)
outputs = model(inputs)
sbp_outputs = outputs[:, :sbp_num_class]
dbp_outputs = outputs[:, sbp_num_class:]
_, sbp_pred = torch.max(sbp_outputs, 1)
_, dbp_pred = torch.max(dbp_outputs, 1)
val_total += inputs.size(0)
val_sbp_correct += (
targets[:, 0] == sbp_pred).sum().item()
val_dbp_correct += (
targets[:, 1] == dbp_pred).sum().item()
sbp_pred_tensor = torch.cat(
(sbp_pred_tensor, sbp_pred), dim=0)
dbp_pred_tensor = torch.cat(
(dbp_pred_tensor, dbp_pred), dim=0)
target_tensor = torch.cat(
(target_tensor, targets.long()), dim=0)
print("\n Acc. on Validation: SBP: {:.3f} DBP:{:.3f}".
format(val_sbp_correct / val_total,
val_dbp_correct / val_total))
writer.add_scalars(
'Acc/Val', {
'SBP': val_sbp_correct / val_total,
'DBP': val_dbp_correct / val_total
}, iteration)
# Save best model
curr_acc = (val_sbp_correct +
val_dbp_correct) / 2 / val_total
if best_acc < curr_acc:
print("Saving best model with acc: {}...".format(
curr_acc))
utils.save_snapshot(model, optimizer,
args.save_result_root, (epoch + 1),
iteration, (epoch + 1))
best_acc = curr_acc
if (epoch + 1) % args.snapshot_epoch_freq == 0 and (
i + 1
) == total_step // args.valid_iter_freq * args.valid_iter_freq:
_, sbp_log = utils.confusion_matrix(
sbp_pred_tensor, target_tensor[:, 0],
sbp_num_class)
_, dbp_log = utils.confusion_matrix(
dbp_pred_tensor, target_tensor[:, 1],
dbp_num_class)
for c in range(sbp_num_class):
per_class = {
'Sensitivity':
sbp_log[0]['class_{}'.format(c)],
'Specificity': sbp_log[1]['class_{}'.format(c)]
}
writer.add_scalars(
'SBP_Metrics/class_{}'.format(c), per_class,
iteration)
for c in range(dbp_num_class):
per_class = {
'Sensitivity':
dbp_log[0]['class_{}'.format(c)],
'Specificity': dbp_log[1]['class_{}'.format(c)]
}
writer.add_scalars(
'DBP_Metrics/class_{}'.format(c), per_class,
iteration)