import tensorflow as tf
import matplotlib .pyplot as plt
from PIL import Image
import numpy as np
import os
from myResNet import resnet35
# 设置相关底层配置
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)
size = 100
mybatch = 30
# l_rate = 5e-7
load_name = 1e-9
l_rate = 1e-9
ep = 500
img_shape = (1,size,size,3)
each_shape = (size,size)
model_shape = (None,size, size, 3)
def getdata():
img_path = 'D:\深度学习\my_imgtest\myimg'
x_shape = img_shape
x_np = np.zeros(shape=x_shape)
y_shape = (1,1)
y_np = np.zeros(shape=y_shape)
y_flag = ['cui','ma','mei','po','yi']
for each_person in os.listdir(img_path):
person_imgs_path = os.path.join(img_path,each_person)
# print(person_imgs_path)
# 确定y标签
flag = person_imgs_path.split('\\')[-1]
flag = flag.split('_')[0]
# print(flag)
flag_index = y_flag.index(flag)+1
# print(flag_index)
for _,_,imgs_name in os.walk(person_imgs_path):
for each_img_name in imgs_name:
each_img = os.path.join(person_imgs_path,each_img_name)
img1 = Image.open(each_img)
img1 = img1.resize(each_shape)
img_np = np.array(img1)
# print(img_np)
# print(img_np.shape)
if np.max(x_np[0]) == 0:
# print('第一个是0')
x_np[0] = img_np
# print('[flag_index]', [flag_index])
# print('y_np:', y_np)
y_np[0] = [flag_index]
continue
x_np = np.insert(x_np,0,img_np,axis=0)
y_np = np.insert(y_np,0,flag_index,axis=0)
# print('x_np:', x_np)
# break
print('读取图片成功！！')
# print('x_np.shape:',x_np.shape)
# print('y_np.shape',y_np.shape)
# x_np.dtype = 'int32'
return x_np,y_np
def preprocess(x,y):
x = tf.cast(x,dtype=tf.float32) / 255
y = tf.cast(y,dtype=tf.int32)
return x,y
if __name__ == '__main__':
x,y = getdata()
# # 画图
# for i in range(x.shape[0]):
# if i != 14 :
# continue
# print(x[i].shape)
# plt.imshow(x[i].astype('int32'))
# y_flag = ['cui', 'ma', 'mei', 'po', 'yi']
# print('y[i]:',y[i][0])
# name = y_flag[int(y[i][0])-1]
# plt.title(name)
# plt.show()
# 识别训练
y = tf.squeeze(y, axis=1)
print(x.shape,y.shape)
train_db = tf.data.Dataset.from_tensor_slices((x, y))
train_db = train_db.map(preprocess).shuffle(30).batch(mybatch)
# 打印看下数据的形状
sample = next(iter(train_db))
print('sample:', sample[0].shape, sample[1].shape
, tf.reduce_min(sample[0]), tf.reduce_max(sample[0]))
model = resnet35()
model.build(input_shape=model_shape)
# 定义优化器
optimizer = tf.optimizers.Adam(lr=l_rate)
# 加载参数
load_path = os.path.join('D:\深度学习\my_imgtest\model_save',str(load_name))
model.load_weights(load_path)
# 训练
for epoch in range(ep):
for step, (x, y) in enumerate(train_db):
with tf.GradientTape() as tape:
# print('x.shape:',x.shape)
logits = model(x)
#
y_onehot = tf.one_hot(y, depth=5)
loss = tf.losses.categorical_crossentropy(y_onehot, logits, from_logits=True)
loss = tf.reduce_mean(loss)
grads = tape.gradient(loss, model.trainable_variables)
if loss > 10:
print('grads:',grads)
break
break
optimizer.apply_gradients(zip(grads, model.trainable_variables))
if step % 100 == 0:
print(epoch, step, 'loss:', float(loss))
for x, y in train_db:
# print()
logits = model(x)
prob = tf.nn.softmax(logits, axis=1)
pred = tf.argmax(prob, axis=1)
pred = tf.cast(pred, tf.int32)
correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
correct = tf.reduce_mean(tf.cast(correct, dtype=tf.float32))
print('acc:', float(correct))
save_path = 'D:\深度学习\my_imgtest\model_save'
save_path = os.path.join(save_path,str(l_rate))
model.save_weights(save_path)
