def evaluate_model(args: argparse.ArgumentParser, learner: Learner):
"""Evaluate the model on the test set."""
learner.load()
test_outputs, test_labels, test_loss = learner.evaluate(
learner.container.test_loader)
test_accuracy = get_accuracy(test_outputs, test_labels)
print(f"test_accuracy: {test_accuracy:.3f}")
def fit_model(args: argparse.ArgumentParser, learner: Learner):
# Fit the model
history = learner.fit(n_epochs=args.epochs)
# Plot training curves
log_dir = f"{learner.container.rootdir}/logs/{learner.model_name}"
plot_training_curves(
"loss",
history,
f"{log_dir}/training_curves/loss.png",
)
plot_training_curves(
"accuracy",
history,
f"{log_dir}/training_curves/accuracy.png",
)
train_outputs, train_labels, train_loss = learner.evaluate(
learner.container.train_loader)
train_accuracy = get_accuracy(train_outputs, train_labels)
print(f"train_accuracy: {train_accuracy:.3f}")
net = SmallPreActResNet(n_classes=container.n_classes)
learner = Learner(container, net, model_name=pruned_model_name)
# learner.get_model_summary()
if args.pruner == "unstructured":
Pruner = UnstructuredPruner
if args.pruner == "structured":
Pruner = StructuredPruner
no_retrain_res = defaultdict(list)
retrain_res = defaultdict(list)
# Load the model
learner.load(model_path)
# Evaluate without retraining
outputs, labels, loss = learner.evaluate(learner.container.test_loader)
accuracy = get_accuracy(outputs, labels)
no_retrain_res[0.0] = [accuracy for _ in range(3)]
retrain_res[0.0] = [accuracy for _ in range(3)]
total_params = sum(p.numel() for p in learner.net.parameters())
log_file.write(
f"[pruning rate: 0.0] n params: {total_params:,} / {accuracy:.2} \n"
)
rates = [0.2, 0.4, 0.6, 0.8]
for current_pruning_rate in rates:
for k_layer in range(1, 4):
# Set the pruning rates for each layer
current_layer = f"dense{k_layer}"
pruning_rates = {f"dense{k}": 0 for k in range(1, 4)}
pruning_rates[current_layer] = current_pruning_rate