def genrate_pre_result():
#获取需要预测结果的所有数据
img_ids, _, img_paths = get_img_infos("test", test_txt)
test_dataset = pd.DataFrame({"img_id": img_ids, "img_path": img_paths})
with tf.device("/cpu:0"):
test_data = ImageDataGenerator(test_dataset,
mode="test",
batch_size=batch_size,
num_classes=num_classes)
iterator = tf.data.Iterator.from_structure(
test_data.data.output_types, test_data.data.output_shapes)
next_batch = iterator.get_next()
#初始化测试集中的图片数据
test_init_op = iterator.make_initializer(test_data.data)
#创建一个加载模型文件的对象
saver = tf.train.Saver()
#用来保存图片的id
test_img_ids = []
#用来保存图片的预测结果
test_pred_labels = []
#计算需要迭代的次数
steps = (test_data.data_size - 1) // batch_size + 1
#设置模型文件的路径
model_path = "checkpoints/model_epoch37_0.9275.ckpt"
print('加载完成')
with tf.Session() as sess:
sess.run(test_init_op)
#加载模型文件
saver.restore(sess, model_path)
for step in range(steps):
#获取数据
image_data, image_id = sess.run(next_batch)
#预测图片的标签
pred_label = sess.run(output_y,
feed_dict={
x: image_data,
keep_prob: 1.0
})
pred_prob = tf.nn.softmax(pred_label)
#保存预测的结果
test_img_ids.extend(image_id)
test_pred_labels.extend(
np.round(sess.run(pred_prob)[:, 1], decimals=2))
data = pd.DataFrame({"id": test_img_ids, "label": test_pred_labels})
data.sort_values(by="id", ascending=True, inplace=True)
#保存结果
data.to_csv("AlexNet_transfer2.csv", index=False)
def train():
label_name_to_num = {"dog": 1, "cat": 0}
#获取所有的训练数据
img_ids, img_labels, img_paths = get_img_infos("train", train_txt,
label_name_to_num)
train_dataset, val_dataset = split_dataset(img_ids, img_paths, img_labels)
print(train_dataset)
with tf.device("/cpu:0"):
train_data = ImageDataGenerator(train_dataset,
mode="train",
batch_size=batch_size,
num_classes=num_classes,
shuffle=True)
val_data = ImageDataGenerator(val_dataset,
mode="val",
batch_size=batch_size,
num_classes=num_classes)
#创建一个获取下一个batch的迭代器
iterator = tf.data.Iterator.from_structure(
train_data.data.output_types, train_data.data.output_shapes)
next_batch = iterator.get_next()
#初始化训练集数据
training_init_op = iterator.make_initializer(train_data.data)
#初始化测试集数据
val_init_op = iterator.make_initializer(val_data.data)
#获取需要重新训练的变量
var_list = [
v for v in tf.trainable_variables()
if v.name.split("/")[0] in train_layers
]
#定义交叉熵损失值
with tf.name_scope("cross_entropy_loss"):
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=output_y, labels=y))
#更新变量
with tf.name_scope("train"):
#计算需要更新变量的梯度
gradients = tf.gradients(loss, var_list)
gradients = list(zip(gradients, var_list))
#更新权重
# optimizer = tf.train.GradientDescentOptimizer(learning_rate)
# train_op = optimizer.apply_gradients(grads_and_vars=gradients)
train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)
for gradient, var in gradients:
tf.summary.histogram(var.name + "/gradient", gradient)
for var in var_list:
tf.summary.histogram(var.name, var)
tf.summary.scalar("cross_entropy", loss)
#计算准确率
with tf.name_scope("accuracy"):
correct_pred = tf.equal(tf.argmax(output_y, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
tf.summary.scalar("accuracy", accuracy)
merged_summary = tf.summary.merge_all()
writer = tf.summary.FileWriter(filewrite_path)
saver = tf.train.Saver()
#计算每轮的迭代次数
train_batches_per_epoch = int(np.floor(train_data.data_size / batch_size))
val_batches_per_epoch = int(np.floor(val_data.data_size / batch_size))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
writer.add_graph(sess.graph)
model.load_initial_weights(sess)
#记录最好的验证准确率
best_val_acc = 0.9
print('开始训练')
#迭代训练
for epoch in range(num_epochs):
sess.run(training_init_op)
for step in range(train_batches_per_epoch):
# print('0batch')
img_batch, label_batch = sess.run(next_batch)
sess.run(train_op,
feed_dict={
x: img_batch,
y: label_batch,
keep_prob: dropout_rate
})
if step % display_step == 0:
s, train_acc, train_loss = sess.run(
[merged_summary, accuracy, loss],
feed_dict={
x: img_batch,
y: label_batch,
keep_prob: 1.0
})
writer.add_summary(s,
epoch * train_batches_per_epoch + step)
sess.run(val_init_op)
#统计验证集的准确率
val_acc = 0
#统计验证集的损失值
val_loss = 0
test_count = 0
for _ in range(val_batches_per_epoch):
img_batch, label_batch = sess.run(next_batch)
acc, val_batch_loss = sess.run([accuracy, loss],
feed_dict={
x: img_batch,
y: label_batch,
keep_prob: 1.0
})
val_acc += acc
val_loss += val_batch_loss
test_count += 1
val_acc /= test_count
val_loss /= test_count
print(
"%s epoch:%d,train acc:%.4f,train loss:%.4f,val acc:%.4f,val loss:%.4f"
% (datetime.now(), epoch + 1, train_acc, train_loss, val_acc,
val_loss))
if val_acc > best_val_acc:
checkpoint_name = os.path.join(
checkpoint_path,
"model_epoch%s_%.4f.ckpt" % (str(epoch + 1), val_acc))
saver.save(sess, checkpoint_name)
best_val_acc = val_acc
for i in range(5):
plt.subplot(2, 4, i + 1)
plt.imshow(Image.open(correct_img_paths[i]))
plt.title("predition:%s\nreal:%s" %
(label_num_to_name[correct_pred_labels[i]],
label_num_to_name[correct_real_labels[i]]))
plt.subplot(2, 5, i + 6)
plt.imshow(Image.open(error_img_paths[i]))
plt.title("prediction:%s\nreal:%s" %
(label_num_to_name[error_pred_labels[i]],
label_num_to_name[error_real_labels[i]]))
plt.show()
#绘制混淆矩阵,输出分类情况报告
def classifiction_report_info(pred_labels, real_labels):
#输出分类结果报告
classification_info = classification_report(
real_labels, pred_labels, target_names=["hurt", "mildew", "worm"])
print(classification_info)
#展示混淆矩阵
con_matrix = confusion_matrix(real_labels, pred_labels)
sns.heatmap(con_matrix, annot=True)
plt.show()
if __name__ == "__main__":
img_ids, img_labels, img_paths = get_img_infos("train", "txt/train.txt")
train_dataset, val_dataset = split_dataset(img_ids, img_paths, img_labels)
plot_label_distribution(train_dataset["img_label"])
plot_label_distribution(val_dataset["img_label"])
default='attention',
help='Choose b/w attention and gated_attention')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
print('\nGPU is ON!')
print('Load Train and Test Set')
loader_kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
# data_utils.DataLoader: 数据加载
img_name, img_path, img_label = util_data.get_img_infos("./data/image.txt")
train_bag_name, test_bag_name = util_data.split_train_test(img_name, img_label)
print('Init Model')
if args.model == 'attention':
model = Attention()
elif args.model == 'gated_attention':
model = GatedAttention()
if args.cuda:
model.cuda()
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
weight_decay=args.reg)
