def test(model):
criterion = loss
test_loss, test_acc, test_r_sq = evaluate(test_loader, model, args.model, criterion)
if args.model == 'popularity':
print('Test Loss: {:.6f} \t'
'Test RSQ.: {:.6f}'.format(test_loss, test_r_sq))
else:
print('Test Loss: {:.6f} \t'
'Test Acc.: {:.6f}'.format(test_loss, test_acc))
# Load best model and test
if args.model == 'artist':
msd = load('artist.pth')
model.load_state_dict(msd)
elif args.model == 'key':
msd = load('key.pth')
model.load_state_dict(msd)
test(model)
# Plot results
plt.plot(range(args.epochs), losses, 'k', range(args.epochs), val_losses, 'b')
plt.legend(['loss', 'val loss'])
plt.show()
if args.model == 'popularity':
plt.plot(range(args.epochs), r_sqs, 'k', range(args.epochs), val_r_sqs, 'b')
plt.legend(['r2', 'val r2'])
else:
plt.plot(range(args.epochs), accs, 'k', range(args.epochs), val_accs, 'b')
