def train_on_dataset(
train_dataset: Dataset, val_dataset, model: Tree2Seq, criterion: nn.modules.loss, optimizer: torch.optim,
scheduler: torch.optim.lr_scheduler, clip_norm: int, logger: AbstractLogger, start_batch_id: int = 0,
log_step: int = -1, eval_step: int = -1, save_step: int = -1
):
train_epoch_info = LearningInfo()
batch_iterator_pb = tqdm(range(start_batch_id, len(train_dataset)), total=len(train_dataset))
batch_iterator_pb.update(start_batch_id)
batch_iterator_pb.refresh()
for batch_id in batch_iterator_pb:
graph, labels = train_dataset[batch_id]
batch_info = train_on_batch(model, criterion, optimizer, scheduler, graph, labels, clip_norm)
train_epoch_info.accumulate_info(batch_info)
if is_step_match(batch_id, log_step):
logger.log(train_epoch_info.get_state_dict(), batch_id, is_train=True)
train_epoch_info = LearningInfo()
if is_step_match(batch_id, save_step):
train_dump = {
'state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'batch_id': batch_id
}
logger.save_model(f'batch_{batch_id}.pt', train_dump)
if is_step_match(batch_id, eval_step):
eval_info = evaluate_on_dataset(val_dataset, model, criterion)
logger.log(eval_info.get_state_dict(), batch_id, is_train=False)
if train_epoch_info.batch_processed > 0:
logger.log(train_epoch_info.get_state_dict(), len(train_dataset) - 1, is_train=True)
def save(epoch: int, model, optimizer: torch.optim.Optimizer,
scheduler: torch.optim.lr_scheduler, config):
"""
Pickles the models to hdd
"""
now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
out_dir = config.output_dir
save_name = os.path.join(out_dir, 'epoch_{}_{}.pth'.format(epoch, now))
save_dict = {
'epoch': epoch,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
}
torch.save(save_dict, save_name)
print("Saved the model to hdd")
def train_helper(model: torchvision.models.resnet.ResNet,
dataloaders: Dict[str, torch.utils.data.DataLoader],
dataset_sizes: Dict[str,
int], criterion: torch.nn.modules.loss,
optimizer: torch.optim, scheduler: torch.optim.lr_scheduler,
num_epochs: int, writer: IO, train_order_writer: IO,
device: torch.device, start_epoch: int, batch_size: int,
save_interval: int, checkpoints_folder: Path, num_layers: int,
classes: List[str], num_classes: int) -> None:
"""
Function for training ResNet.
Args:
model: ResNet model for training.
dataloaders: Dataloaders for IO pipeline.
dataset_sizes: Sizes of the training and validation dataset.
criterion: Metric used for calculating loss.
optimizer: Optimizer to use for gradient descent.
scheduler: Scheduler to use for learning rate decay.
start_epoch: Starting epoch for training.
writer: Writer to write logging information.
train_order_writer: Writer to write the order of training examples.
device: Device to use for running model.
num_epochs: Total number of epochs to train for.
batch_size: Mini-batch size to use for training.
save_interval: Number of epochs between saving checkpoints.
checkpoints_folder: Directory to save model checkpoints to.
num_layers: Number of layers to use in the ResNet model from [18, 34, 50, 101, 152].
classes: Names of the classes in the dataset.
num_classes: Number of classes in the dataset.
"""
since = time.time()
# Initialize all the tensors to be used in training and validation.
# Do this outside the loop since it will be written over entirely at each
# epoch and doesn't need to be reallocated each time.
train_all_labels = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
train_all_predicts = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
val_all_labels = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
val_all_predicts = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
global_minibatch_counter = 0
# Train for specified number of epochs.
for epoch in range(start_epoch, num_epochs):
# Training phase.
model.train(mode=True)
train_running_loss = 0.0
train_running_corrects = 0
epoch_minibatch_counter = 0
# Train over all training data.
for idx, (inputs, labels, paths) in enumerate(dataloaders["train"]):
train_inputs = inputs.to(device=device)
train_labels = labels.to(device=device)
optimizer.zero_grad()
# Forward and backpropagation.
with torch.set_grad_enabled(mode=True):
train_outputs = model(train_inputs)
__, train_preds = torch.max(train_outputs, dim=1)
train_loss = criterion(input=train_outputs,
target=train_labels)
train_loss.backward()
optimizer.step()
# Update training diagnostics.
train_running_loss += train_loss.item() * train_inputs.size(0)
train_running_corrects += torch.sum(
train_preds == train_labels.data, dtype=torch.double)
start = idx * batch_size
end = start + batch_size
train_all_labels[start:end] = train_labels.detach().cpu()
train_all_predicts[start:end] = train_preds.detach().cpu()
global_minibatch_counter += 1
epoch_minibatch_counter += 1
# for path in paths: #write the order that the model was trained in
# train_order_writer.write("/".join(path.split("/")[-2:]) + "\n")
if global_minibatch_counter % 10 == 0 or global_minibatch_counter == 5:
calculate_confusion_matrix(
all_labels=train_all_labels.numpy(),
all_predicts=train_all_predicts.numpy(),
classes=classes,
num_classes=num_classes)
# Store training diagnostics.
train_loss = train_running_loss / (epoch_minibatch_counter *
batch_size)
train_acc = train_running_corrects / (epoch_minibatch_counter *
batch_size)
# Validation phase.
model.train(mode=False)
val_running_loss = 0.0
val_running_corrects = 0
# Feed forward over all the validation data.
for idx, (val_inputs, val_labels,
paths) in enumerate(dataloaders["val"]):
val_inputs = val_inputs.to(device=device)
val_labels = val_labels.to(device=device)
# Feed forward.
with torch.set_grad_enabled(mode=False):
val_outputs = model(val_inputs)
_, val_preds = torch.max(val_outputs, dim=1)
val_loss = criterion(input=val_outputs,
target=val_labels)
# Update validation diagnostics.
val_running_loss += val_loss.item() * val_inputs.size(0)
val_running_corrects += torch.sum(
val_preds == val_labels.data, dtype=torch.double)
start = idx * batch_size
end = start + batch_size
val_all_labels[start:end] = val_labels.detach().cpu()
val_all_predicts[start:end] = val_preds.detach().cpu()
calculate_confusion_matrix(
all_labels=val_all_labels.numpy(),
all_predicts=val_all_predicts.numpy(),
classes=classes,
num_classes=num_classes)
# Store validation diagnostics.
val_loss = val_running_loss / dataset_sizes["val"]
val_acc = val_running_corrects / dataset_sizes["val"]
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Remaining things related to training.
if global_minibatch_counter % 10 == 0 or global_minibatch_counter == 5:
epoch_output_path = checkpoints_folder.joinpath(
f"resnet{num_layers}_e{epoch}_mb{global_minibatch_counter}_va{val_acc:.5f}.pt"
)
# Confirm the output directory exists.
epoch_output_path.parent.mkdir(parents=True, exist_ok=True)
# Save the model as a state dictionary.
torch.save(obj={
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"epoch": epoch + 1
},
f=str(epoch_output_path))
writer.write(
f"{epoch},{global_minibatch_counter},{train_loss:.4f},"
f"{train_acc:.4f},{val_loss:.4f},{val_acc:.4f}\n")
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print the diagnostics for each epoch.
print(f"Epoch {epoch} with "
f"mb {global_minibatch_counter} "
f"lr {current_lr:.15f}: "
f"t_loss: {train_loss:.4f} "
f"t_acc: {train_acc:.4f} "
f"v_loss: {val_loss:.4f} "
f"v_acc: {val_acc:.4f}\n")
scheduler.step()
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print training information at the end.
print(f"\ntraining complete in "
f"{(time.time() - since) // 60:.2f} minutes")
def save_checkpoint(
path: str,
model: TEDD1104,
optimizer_name: str,
optimizer: torch.optim,
scheduler: torch.optim.lr_scheduler,
acc_dev: float,
epoch: int,
fp16: bool,
opt_level: str = None,
) -> None:
"""
Save a checkpoint that allows to continue training the model in the future
Input:
- path: path where the model is going to be saved
- model: TEDD1104 model to save
- optimizer_name: Name of the optimizer used for training: SGD or Adam
- optimizer: Optimizer used for training
- acc_dev: Accuracy of the model in the development set
- epoch: Num of epoch used to train the model
- amp: If the model uses FP16, Nvidia Apex AMP
- amp_opt_level: If the model uses FP16, the AMP opt_level
Output:
"""
if fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
dict_hyperparams: dict = {
"sequence_size": model.sequence_size,
"resnet": model.resnet,
"pretrained_resnet": model.pretrained_resnet,
"embedded_size": model.embedded_size,
"hidden_size": model.hidden_size,
"num_layers_lstm": model.num_layers_lstm,
"bidirectional_lstm": model.bidirectional_lstm,
"layers_out": model.layers_out,
"dropout_cnn": model.dropout_cnn,
"dropout_cnn_out": model.dropout_cnn_out,
"dropout_lstm": model.dropout_lstm,
"dropout_lstm_out": model.dropout_lstm_out,
"fp16": fp16,
"amp_opt_level": opt_level,
}
checkpoint = {
"hyper_params": dict_hyperparams,
"model": model.state_dict(),
"optimizer_name": optimizer_name,
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"acc_dev": acc_dev,
"epoch": epoch,
"amp": None if not fp16 else amp.state_dict(),
"opt_level": opt_level,
}
torch.save(checkpoint, path)
def train_helper_with_gradients(model: torchvision.models.resnet.ResNet,
dataloaders: Dict[str,
torch.utils.data.DataLoader],
dataset_sizes: Dict[str, int],
criterion: torch.nn.modules.loss,
optimizer: torch.optim,
scheduler: torch.optim.lr_scheduler,
num_epochs: int, writer: IO,
train_order_writer: IO, device: torch.device,
start_epoch: int, batch_size: int,
save_interval: int, checkpoints_folder: Path,
num_layers: int, classes: List[str],
num_classes: int) -> None:
since = time.time()
# Initialize all the tensors to be used in training and validation.
# Do this outside the loop since it will be written over entirely at each
# epoch and doesn't need to be reallocated each time.
train_all_labels = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
train_all_predicts = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
val_all_labels = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
val_all_predicts = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
global_minibatch_counter = 0
mag_writer = open("mags_resnet18_imagenet.csv", "w")
mag_writer.write(
"image_name,train_loss,layers_-1,layer_0,layer_60,layer_1,layer_20,layer_40,layer_59,conf,correct\n"
)
# Train for specified number of epochs.
for epoch in range(start_epoch, num_epochs):
# Training phase.
model.train(mode=True)
train_running_loss = 0.0
train_running_corrects = 0
epoch_minibatch_counter = 0
# Train over all training data.
for idx, (inputs, labels, paths) in enumerate(dataloaders["train"]):
train_inputs = inputs.to(device=device)
train_labels = labels.to(device=device)
optimizer.zero_grad()
# Forward and backpropagation.
with torch.set_grad_enabled(mode=True):
train_outputs = model(train_inputs)
confs, train_preds = torch.max(train_outputs, dim=1)
train_loss = criterion(input=train_outputs,
target=train_labels)
train_loss.backward(retain_graph=True)
optimizer.step()
batch_grads = torch.autograd.grad(train_loss,
model.parameters(),
retain_graph=True)
# print(len(batch_grads))
# for batch_grad in batch_grads:
# print(batch_grad.size())
train_loss_npy = float(train_loss.detach().cpu().numpy())
layer_num_to_mag = get_grad_magnitude(model)
image_name = get_image_name(paths[0])
conf = float(confs.detach().cpu().numpy())
train_pred = int(train_preds.detach().cpu().numpy()[0])
gt_label = int(train_labels.detach().cpu().numpy()[0])
correct = 0
if train_pred == gt_label:
correct = 1
output_line = f"{image_name},{train_loss_npy:.4f},{layer_num_to_mag[-1]:.4f},{layer_num_to_mag[0]:.4f},{layer_num_to_mag[60]:.4f},{layer_num_to_mag[1]:.4f},{layer_num_to_mag[20]:.4f},{layer_num_to_mag[40]:.4f},{layer_num_to_mag[59]:.4f},{conf:.4f},{correct}\n"
mag_writer.write(output_line)
print(idx, output_line)
# print(idx, image_name, train_loss_npy, conf, train_pred, gt_label)
# Update training diagnostics.
train_running_loss += train_loss.item() * train_inputs.size(0)
train_running_corrects += torch.sum(
train_preds == train_labels.data, dtype=torch.double)
start = idx * batch_size
end = start + batch_size
train_all_labels[start:end] = train_labels.detach().cpu()
train_all_predicts[start:end] = train_preds.detach().cpu()
global_minibatch_counter += 1
epoch_minibatch_counter += 1
if global_minibatch_counter % 1000 == 0:
calculate_confusion_matrix(
all_labels=train_all_labels.numpy(),
all_predicts=train_all_predicts.numpy(),
classes=classes,
num_classes=num_classes)
# Store training diagnostics.
train_loss = train_running_loss / (epoch_minibatch_counter *
batch_size)
train_acc = train_running_corrects / (epoch_minibatch_counter *
batch_size)
# Validation phase.
model.train(mode=False)
val_running_loss = 0.0
val_running_corrects = 0
# Feed forward over all the validation data.
for idx, (val_inputs, val_labels,
paths) in enumerate(dataloaders["val"]):
val_inputs = val_inputs.to(device=device)
val_labels = val_labels.to(device=device)
# Feed forward.
with torch.set_grad_enabled(mode=False):
val_outputs = model(val_inputs)
_, val_preds = torch.max(val_outputs, dim=1)
val_loss = criterion(input=val_outputs,
target=val_labels)
# Update validation diagnostics.
val_running_loss += val_loss.item() * val_inputs.size(0)
val_running_corrects += torch.sum(
val_preds == val_labels.data, dtype=torch.double)
start = idx * batch_size
end = start + batch_size
val_all_labels[start:end] = val_labels.detach().cpu()
val_all_predicts[start:end] = val_preds.detach().cpu()
calculate_confusion_matrix(
all_labels=val_all_labels.numpy(),
all_predicts=val_all_predicts.numpy(),
classes=classes,
num_classes=num_classes)
# Store validation diagnostics.
val_loss = val_running_loss / dataset_sizes["val"]
val_acc = val_running_corrects / dataset_sizes["val"]
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Remaining things related to training.
if global_minibatch_counter % 200000 == 0 or global_minibatch_counter == 5:
epoch_output_path = checkpoints_folder.joinpath(
f"resnet{num_layers}_e{epoch}_mb{global_minibatch_counter}_va{val_acc:.5f}.pt"
)
# Confirm the output directory exists.
epoch_output_path.parent.mkdir(parents=True, exist_ok=True)
# Save the model as a state dictionary.
torch.save(obj={
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"epoch": epoch + 1
},
f=str(epoch_output_path))
writer.write(
f"{epoch},{global_minibatch_counter},{train_loss:.4f},"
f"{train_acc:.4f},{val_loss:.4f},{val_acc:.4f}\n")
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print the diagnostics for each epoch.
print(f"Epoch {epoch} with "
f"mb {global_minibatch_counter} "
f"lr {current_lr:.15f}: "
f"t_loss: {train_loss:.4f} "
f"t_acc: {train_acc:.4f} "
f"v_loss: {val_loss:.4f} "
f"v_acc: {val_acc:.4f}\n")
scheduler.step()
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print training information at the end.
print(f"\ntraining complete in "
f"{(time.time() - since) // 60:.2f} minutes")
def save_checkpoint(
path: str,
model: TEDD1104,
optimizer_name: str,
optimizer: torch.optim,
scheduler: torch.optim.lr_scheduler,
running_loss: float,
total_batches: int,
total_training_examples: int,
acc_dev: float,
epoch: int,
fp16: bool,
scaler: Optional[GradScaler],
) -> None:
"""
Save a checkpoint that allows to continue training the model in the future
Input:
- path: path where the model is going to be saved
- model: TEDD1104 model to save
- optimizer_name: Name of the optimizer used for training: SGD or Adam
- optimizer: Optimizer used for training
- acc_dev: Accuracy of the model in the development set
- epoch: Num of epoch used to train the model
- fp16: If the model uses FP16
- scaler: If the model uses FP16, the scaler used for training
Output:
"""
dict_hyperparams: dict = {
"sequence_size": model.sequence_size,
"resnet": model.resnet,
"pretrained_resnet": model.pretrained_resnet,
"embedded_size": model.embedded_size,
"hidden_size": model.hidden_size,
"num_layers_lstm": model.num_layers_lstm,
"bidirectional_lstm": model.bidirectional_lstm,
"layers_out": model.layers_out,
"dropout_cnn": model.dropout_cnn,
"dropout_cnn_out": model.dropout_cnn_out,
"dropout_lstm": model.dropout_lstm,
"dropout_lstm_out": model.dropout_lstm_out,
"fp16": fp16,
}
checkpoint = {
"hyper_params": dict_hyperparams,
"model": model.state_dict(),
"optimizer_name": optimizer_name,
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"running_loss": running_loss,
"total_batches": total_batches,
"total_training_examples": total_training_examples,
"acc_dev": acc_dev,
"epoch": epoch,
"scaler": None if not fp16 else scaler.state_dict(),
}
torch.save(checkpoint, path)
def _train_helper(self, model: torchvision.models.resnet.ResNet,
dataloaders: Dict[str, torch.utils.data.DataLoader],
dataset_sizes: Dict[str, int], loss_fn,
optimizer: torch.optim,
scheduler: torch.optim.lr_scheduler, start_epoch: int,
writer: IO) -> None:
"""
Function for learning ResNet.
Args:
model: ResNet model for learning.
dataloaders: Dataloaders for IO pipeline.
dataset_sizes: Sizes of the learning and validation dataset.
loss_fn: Metric used for calculating loss.
optimizer: Optimizer to use for gradient descent.
scheduler: Scheduler to use for learning rate decay.
start_epoch: Starting epoch for learning.
writer: Writer to write logging information.
"""
learning_init_time = time.time()
# Initialize all the tensors to be used in learning and validation.
# Do this outside the loop since it will be written over entirely at each
# epoch and doesn't need to be reallocated each time.
train_all_labels = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
train_all_predicts = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
val_all_labels = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
val_all_predicts = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
early_stopper = EarlyStopper(patience=self._early_stopping_patience,
mode=EarlyStopper.Mode.MAX)
if self._resume_checkpoint and self._last_val_acc:
best_val_acc = self._last_val_acc
else:
best_val_acc = 0.
# Train for specified number of epochs.
for epoch in range(start_epoch, self._num_epochs):
epoch_init_time = time.time()
# Training phase.
model.train(mode=True)
train_running_loss = 0.0
train_running_corrects = 0
# Train over all learning data.
for idx, (train_inputs,
true_labels) in enumerate(dataloaders["train"]):
train_patches = train_inputs["patch"].to(device=self._device)
train_x_coord = train_inputs["x_coord"].to(device=self._device)
train_y_coord = train_inputs["y_coord"].to(device=self._device)
true_labels = true_labels.to(device=self._device)
optimizer.zero_grad()
# Forward and backpropagation.
with torch.set_grad_enabled(mode=True):
train_logits = model(train_patches, train_x_coord,
train_y_coord).squeeze(dim=1)
train_loss = loss_fn(logits=train_logits,
target=true_labels)
train_loss.backward()
optimizer.step()
# Update learning diagnostics.
train_running_loss += train_loss.item() * train_patches.size(0)
pred_labels = self._extract_pred_labels(train_logits)
train_running_corrects += torch.sum(
pred_labels == true_labels.data, dtype=torch.double)
start = idx * self._batch_size
end = start + self._batch_size
train_all_labels[start:end] = true_labels.detach().cpu()
train_all_predicts[start:end] = pred_labels.detach().cpu()
self._calculate_confusion_matrix(
all_labels=train_all_labels.numpy(),
all_predicts=train_all_predicts.numpy(),
classes=self._classes,
num_classes=self._num_classes)
# Store learning diagnostics.
train_loss = train_running_loss / dataset_sizes["train"]
train_acc = train_running_corrects / dataset_sizes["train"]
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Validation phase.
model.train(mode=False)
val_running_loss = 0.0
val_running_corrects = 0
# Feed forward over all the validation data.
for idx, (val_inputs, val_labels) in enumerate(dataloaders["val"]):
val_patches = val_inputs["patch"].to(device=self._device)
val_x_coord = val_inputs["x_coord"].to(device=self._device)
val_y_coord = val_inputs["y_coord"].to(device=self._device)
val_labels = val_labels.to(device=self._device)
# Feed forward.
with torch.set_grad_enabled(mode=False):
val_logits = model(val_patches, val_x_coord,
val_y_coord).squeeze(dim=1)
val_loss = loss_fn(logits=val_logits, target=val_labels)
# Update validation diagnostics.
val_running_loss += val_loss.item() * val_patches.size(0)
pred_labels = self._extract_pred_labels(val_logits)
val_running_corrects += torch.sum(
pred_labels == val_labels.data, dtype=torch.double)
start = idx * self._batch_size
end = start + self._batch_size
val_all_labels[start:end] = val_labels.detach().cpu()
val_all_predicts[start:end] = pred_labels.detach().cpu()
self._calculate_confusion_matrix(
all_labels=val_all_labels.numpy(),
all_predicts=val_all_predicts.numpy(),
classes=self._classes,
num_classes=self._num_classes)
# Store validation diagnostics.
val_loss = val_running_loss / dataset_sizes["val"]
val_acc = val_running_corrects / dataset_sizes["val"]
if torch.cuda.is_available():
torch.cuda.empty_cache()
scheduler.step()
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Remaining things related to learning.
if val_acc > best_val_acc:
best_val_acc = val_acc
best_model_ckpt_path = self._checkpoints_folder.joinpath(
f"resnet{self._num_layers}_e{epoch}_va{val_acc:.5f}.pt")
# Confirm the output directory exists.
best_model_ckpt_path.parent.mkdir(parents=True, exist_ok=True)
# Save the model as a state dictionary.
torch.save(obj={
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"epoch": epoch + 1
},
f=str(best_model_ckpt_path))
self._clean_ckpt_folder(best_model_ckpt_path)
writer.write(f"{epoch},{train_loss:.4f},"
f"{train_acc:.4f},{val_loss:.4f},{val_acc:.4f}\n")
# Print the diagnostics for each epoch.
logging.info(
f"Epoch {epoch} "
f"with lr {current_lr:.15f}: "
f"{self._format_time_period(epoch_init_time, time.time())} "
f"t_loss: {train_loss:.4f} "
f"t_acc: {train_acc:.4f} "
f"v_loss: {val_loss:.4f} "
f"v_acc: {val_acc:.4f}\n")
early_stopper.update(val_acc)
if early_stopper.is_stopping():
logging.info("Early stopping")
break
# Print learning information at the end.
logging.info(
f"\nlearning complete in "
f"{self._format_time_period(learning_init_time, time.time())}")
def train_helper(model: torchvision.models.resnet.ResNet,
dataloaders: Dict[str, torch.utils.data.DataLoader],
dataset_sizes: Dict[str,
int], criterion: torch.nn.modules.loss,
optimizer: torch.optim, scheduler: torch.optim.lr_scheduler,
num_epochs: int, log_writer: IO, train_order_writer: IO,
device: torch.device, batch_size: int,
checkpoints_folder: Path, num_layers: int, classes: List[str],
minibatch_counter, num_classes: int) -> None:
since = time.time()
global_minibatch_counter = minibatch_counter
# Initialize all the tensors to be used in training and validation.
# Do this outside the loop since it will be written over entirely at each
# epoch and doesn't need to be reallocated each time.
train_all_labels = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
train_all_predicts = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
val_all_labels = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
val_all_predicts = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
for epoch in range(1, num_epochs + 1):
model.train(mode=True) # Training phase.
train_running_loss, train_running_corrects, epoch_minibatch_counter = 0.0, 0, 0
for idx, (inputs, labels, paths) in enumerate(dataloaders["train"]):
train_inputs = inputs.to(device=device)
train_labels = labels.to(device=device)
optimizer.zero_grad()
# Forward and backpropagation.
with torch.set_grad_enabled(mode=True):
train_outputs = model(train_inputs)
__, train_preds = torch.max(train_outputs, dim=1)
train_loss = criterion(input=train_outputs,
target=train_labels)
train_loss.backward()
optimizer.step()
# Update training diagnostics.
train_running_loss += train_loss.item() * train_inputs.size(0)
train_running_corrects += torch.sum(
train_preds == train_labels.data, dtype=torch.double)
this_batch_size = train_labels.detach().cpu().shape[0]
start = idx * batch_size
end = start + this_batch_size
train_all_labels[start:end] = train_labels.detach().cpu()
train_all_predicts[start:end] = train_preds.detach().cpu()
global_minibatch_counter += 1
epoch_minibatch_counter += 1
# Calculate training diagnostics
calculate_confusion_matrix(all_labels=train_all_labels.numpy(),
all_predicts=train_all_predicts.numpy(),
classes=classes,
num_classes=num_classes)
train_loss = train_running_loss / (epoch_minibatch_counter *
batch_size)
train_acc = train_running_corrects / (epoch_minibatch_counter *
batch_size)
# Validation phase.
model.train(mode=False)
val_running_loss = 0.0
val_running_corrects = 0
# Feed forward over all the validation data.
for idx, (val_inputs, val_labels,
paths) in enumerate(dataloaders["val"]):
val_inputs = val_inputs.to(device=device)
val_labels = val_labels.to(device=device)
# Feed forward.
with torch.set_grad_enabled(mode=False):
val_outputs = model(val_inputs)
_, val_preds = torch.max(val_outputs, dim=1)
val_loss = criterion(input=val_outputs, target=val_labels)
# Update validation diagnostics.
val_running_loss += val_loss.item() * val_inputs.size(0)
val_running_corrects += torch.sum(val_preds == val_labels.data,
dtype=torch.double)
this_batch_size = val_labels.detach().cpu().shape[0]
start = idx * batch_size
end = start + this_batch_size
val_all_labels[start:end] = val_labels.detach().cpu()
val_all_predicts[start:end] = val_preds.detach().cpu()
# Calculate validation diagnostics
calculate_confusion_matrix(all_labels=val_all_labels.numpy(),
all_predicts=val_all_predicts.numpy(),
classes=classes,
num_classes=num_classes)
val_loss = val_running_loss / dataset_sizes["val"]
val_acc = val_running_corrects / dataset_sizes["val"]
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Remaining things related to training.
epoch_output_path = checkpoints_folder.joinpath(
f"resnet{num_layers}_e{epoch}_mb{global_minibatch_counter}_va{val_acc:.5f}.pt"
)
epoch_output_path.parent.mkdir(parents=True, exist_ok=True)
# Save the model as a state dictionary.
torch.save(obj={
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"epoch": epoch + 1
},
f=str(epoch_output_path))
log_writer.write(
f"{epoch},{global_minibatch_counter},{train_loss:.4f},{train_acc:.4f},{val_loss:.4f},{val_acc:.4f}\n"
)
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print the diagnostics for each epoch.
print(f"Epoch {epoch} with "
f"mb {global_minibatch_counter} "
f"lr {current_lr:.15f}: "
f"t_loss: {train_loss:.4f} "
f"t_acc: {train_acc:.4f} "
f"v_loss: {val_loss:.4f} "
f"v_acc: {val_acc:.4f}\n")
scheduler.step()
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print training information at the end.
print(f"\ntraining complete in "
f"{(time.time() - since) // 60:.2f} minutes")
return epoch_output_path, global_minibatch_counter
def train_smartgrad_helper(model: torchvision.models.resnet.ResNet,
dataloaders: Dict[str, torch.utils.data.DataLoader],
dataset_sizes: Dict[str, int],
criterion: torch.nn.modules.loss,
optimizer: torch.optim,
scheduler: torch.optim.lr_scheduler,
num_epochs: int,
log_writer: IO,
train_order_writer: IO,
device: torch.device,
train_batch_size: int,
val_batch_size: int,
fake_minibatch_size: int,
annealling_factor: float,
save_mb_interval: int,
val_mb_interval: int,
checkpoints_folder: Path,
num_layers: int,
classes: List[str],
num_classes: int) -> None:
grad_layers = list(range(1, 21))
since = time.time()
global_minibatch_counter = 0
# Initialize all the tensors to be used in training and validation.
# Do this outside the loop since it will be written over entirely at each
# epoch and doesn't need to be reallocated each time.
train_all_labels = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
train_all_predicts = torch.empty(size=(dataset_sizes["train"], ),
dtype=torch.long).cpu()
val_all_labels = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
val_all_predicts = torch.empty(size=(dataset_sizes["val"], ),
dtype=torch.long).cpu()
for epoch in range(1, num_epochs+1):
model.train(mode=False) # Training phase.
train_running_loss, train_running_corrects, epoch_minibatch_counter = 0.0, 0, 0
idx_to_gt = {}
for idx, (inputs, labels, paths) in enumerate(dataloaders["train"]):
train_inputs = inputs.to(device=device)
train_labels = labels.to(device=device)
optimizer.zero_grad()
# Forward and backpropagation.
with torch.set_grad_enabled(mode=True):
train_outputs = model(train_inputs)
__, train_preds = torch.max(train_outputs, dim=1)
train_loss = criterion(input=train_outputs, target=train_labels)
train_loss.backward(retain_graph=True)
gt_label = int(train_labels.detach().cpu().numpy()[0])
idx_to_gt[idx] = gt_label
########################
#### important code ####
########################
#clear the memory
fake_minibatch_idx = idx % fake_minibatch_size
fake_minibatch_num = int(idx / fake_minibatch_size)
if fake_minibatch_idx == 0:
minibatch_grad_dict = {}; gc.collect()
#get the per-example gradient magnitude and add to minibatch_grad_dict
grad_as_dict, grad_flattened = model_to_grad_as_dict_and_flatten(model, grad_layers)
minibatch_grad_dict[idx] = (grad_as_dict, grad_flattened)
#every batch, calculate the best ones
if fake_minibatch_idx == fake_minibatch_size - 1:
idx_to_weight_batch = get_idx_to_weight(minibatch_grad_dict, annealling_factor, idx_to_gt)
print(idx_to_weight_batch)
##########################
# print("\n...............................updating......................................" + str(idx))
for layer_num, param in enumerate(model.parameters()):
# if layer_num in [0]:#grad_layers:
new_grad = get_new_layer_grad(layer_num, idx_to_weight_batch, minibatch_grad_dict)
assert param.grad.detach().cpu().numpy().shape == new_grad.detach().cpu().numpy().shape
param.grad = new_grad
# check_model_weights(idx, model)
optimizer.step()
# check_model_weights(idx, model)
# print("................................done........................................." + str(idx) + '\n\n\n\n')
##########################
# Update training diagnostics.
train_running_loss += train_loss.item() * train_inputs.size(0)
train_running_corrects += torch.sum(train_preds == train_labels.data, dtype=torch.double)
start = idx * train_batch_size
end = start + train_batch_size
train_all_labels[start:end] = train_labels.detach().cpu()
train_all_predicts[start:end] = train_preds.detach().cpu()
global_minibatch_counter += 1
epoch_minibatch_counter += 1
# Write the path of training order if it exists
if train_order_writer:
for path in paths: #write the order that the model was trained in
train_order_writer.write("/".join(path.split("/")[-2:]) + "\n")
# Validate the model
if global_minibatch_counter % val_mb_interval == 0 or global_minibatch_counter == 1:
# Calculate training diagnostics
calculate_confusion_matrix( all_labels=train_all_labels.numpy(), all_predicts=train_all_predicts.numpy(),
classes=classes, num_classes=num_classes)
train_loss = train_running_loss / (epoch_minibatch_counter * train_batch_size)
train_acc = train_running_corrects / (epoch_minibatch_counter * train_batch_size)
# Validation phase.
model.train(mode=False)
val_running_loss = 0.0
val_running_corrects = 0
# Feed forward over all the validation data.
for idx, (val_inputs, val_labels, paths) in enumerate(dataloaders["val"]):
val_inputs = val_inputs.to(device=device)
val_labels = val_labels.to(device=device)
# Feed forward.
with torch.set_grad_enabled(mode=False):
val_outputs = model(val_inputs)
_, val_preds = torch.max(val_outputs, dim=1)
val_loss = criterion(input=val_outputs, target=val_labels)
# Update validation diagnostics.
val_running_loss += val_loss.item() * val_inputs.size(0)
val_running_corrects += torch.sum(val_preds == val_labels.data,
dtype=torch.double)
start = idx * val_batch_size
end = start + val_batch_size
val_all_labels[start:end] = val_labels.detach().cpu()
val_all_predicts[start:end] = val_preds.detach().cpu()
# Calculate validation diagnostics
calculate_confusion_matrix( all_labels=val_all_labels.numpy(), all_predicts=val_all_predicts.numpy(),
classes=classes, num_classes=num_classes)
val_loss = val_running_loss / dataset_sizes["val"]
val_acc = val_running_corrects / dataset_sizes["val"]
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Remaining things related to training.
if global_minibatch_counter % save_mb_interval == 0 or global_minibatch_counter == 1:
epoch_output_path = checkpoints_folder.joinpath(f"resnet{num_layers}_e{epoch}_mb{global_minibatch_counter}_va{val_acc:.5f}.pt")
epoch_output_path.parent.mkdir(parents=True, exist_ok=True)
# Save the model as a state dictionary.
torch.save(obj={
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"epoch": epoch + 1
}, f=str(epoch_output_path))
log_writer.write(f"{epoch},{global_minibatch_counter},{train_loss:.4f},{train_acc:.4f},{val_loss:.4f},{val_acc:.4f}\n")
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
# Print the diagnostics for each epoch.
print(f"Epoch {epoch} with "
f"mb {global_minibatch_counter} "
f"lr {current_lr:.15f}: "
f"t_loss: {train_loss:.4f} "
f"t_acc: {train_acc:.4f} "
f"v_loss: {val_loss:.4f} "
f"v_acc: {val_acc:.4f}\n")
scheduler.step()
current_lr = None
for group in optimizer.param_groups:
current_lr = group["lr"]
