def alexnet_layers(n, MODEL_NAME):
branchyNet = AlexNet().build(n)
branchyNet.to_gpu()
branchyNet.training()
branchyNet.verbose = False
# 在主网络上进行训练
main_loss, main_acc, main_time = utils.train(branchyNet,
X_train,
Y_train,
main=True,
batchsize=TRAIN_BATCHSIZE,
num_epoch=TRAIN_NUM_EPOCHES)
print("main_time:", main_time)
# 在带有分支的网络上进行训练
branch_loss, branch_acc, branch_time = utils.train(
branchyNet,
X_train,
Y_train,
batchsize=TRAIN_BATCHSIZE,
num_epoch=TRAIN_NUM_EPOCHES)
print("branch_time:", branch_time)
# 保存已经训练好的网络模型
with open(MODEL_PATH + MODEL_NAME, "wb") as f:
dill.dump(branchyNet, f)
# 加载已经保存好的模型
with open(MODEL_PATH + MODEL_NAME, "rb") as f:
branchyNet = dill.load(f)
# 网络模型测试，获取分支网络的准确率
branchyNet.testing()
branchyNet.verbose = False
branchyNet.to_gpu()
# 退出点阈值设置
thresholdExits = [0.05, 0.01]
branchyNet.thresholdExits = thresholdExits
# 根据退出点阈值来获取分支网络的总体准确率
overacc, _, _, accbreakdowns, _ = utils.test(branchyNet,
X_test,
Y_test,
batchsize=TEST_BATCHSIZE)
return accbreakdowns
TRAIN_NUM_EPOCHES = 100
SAVE_PATH = '../pic/lenet_mnist/' # 实验结果图片保存路径
MODEL_NAME = '../models/lenet_mnist(' + str(
TRAIN_NUM_EPOCHES) + ').bn' # 保存模型名称
CSV_NAME = 'lenet(' + str(TRAIN_NUM_EPOCHES) + ')' # 输出文件名称
# 导入MNIST数据集
X_train, Y_train, X_test, Y_test = mnist.get_data()
print("X_train:{} Y_train:{}".format(X_train.shape, Y_train.shape))
print("X_test: {} Y_test: {}".format(X_test.shape, Y_test.shape))
# 在主网络上进行训练
main_loss, main_acc, main_time = utils.train(branchyNet,
X_train,
Y_train,
main=True,
batchsize=TRAIN_BATCHSIZE,
num_epoch=TRAIN_NUM_EPOCHES)
print("main_time:", main_time)
# 可视化主网络训练结果
visualize.plot_layers(main_loss,
save_path=SAVE_PATH,
save_name='main_loss(' + str(TRAIN_NUM_EPOCHES) + ')',
xlabel='Epoches',
ylabel='Training Loss')
visualize.plot_layers(main_acc,
save_path=SAVE_PATH,
save_name='main_acc(' + str(TRAIN_NUM_EPOCHES) + ')',
xlabel='Epoches',
ylabel='Training Accuracy')
from networks import alex_cifar10
from datasets import pcifar10
from branchynet import utils
from config import *
if __name__ == '__main__':
#Get the B-AlexNet architecture to classify images using Cifar10 dataset
branchyNet = alex_cifar10.get_network()
dataloader = pcifar10.get_data()
branchyNet.to_gpu()
branchyNet.training()
# Train only the main branch of BranchyNet.
main_loss, main_acc = utils.train(branchyNet,
dataloader,
main=True,
batchsize=TRAIN_BATCHSIZE,
num_epoch=MAIN_TRAIN_NUM_EPOCHS)
# Train the side branches of BranchyNet.
main_loss, main_acc = utils.train(branchyNet,
dataloader,
batchsize=TRAIN_BATCHSIZE,
num_epoch=BRANCH_TRAIN_NUM_EPOCHS)
branchyNet.save_branchyNet('trained_model/BranchyAlexNet(100,200).pt')
