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 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"]