def train(args):
# parse config
place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
with fluid.dygraph.guard(place):
config = parse_config(args.config)
train_config = merge_configs(config, 'train', vars(args))
#根据自己定义的网络，声明train_model
train_model = TSN(args.batch_size, 32)
train_model.train()
opt = fluid.optimizer.Momentum(0.001,
0.9,
parameter_list=train_model.parameters())
if args.pretrain:
# 加载上一次训练的模型，继续训练
model, _ = fluid.dygraph.load_dygraph(args.save_dir + '/tsn_model')
train_model.load_dict(model)
# build model
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# 迭代器，每次输送一个批次的数据，每个数据的形式(imgs, label)
train_reader = UCFReader('train', train_config).create_reader()
epochs = args.epoch
for i in range(epochs):
for batch_id, data in enumerate(train_reader()):
dy_x_data = np.array([x[0] for x in data]).astype(
'float32'
) # dy_x_data=imgs shape:(batch_size, seg_num*seglen, 3, target_size, target_size)
y_data = np.array([x[1] for x in data]).astype(
'int64') # y_data=label shape:(batch_size,)
img = fluid.dygraph.to_variable(dy_x_data)
label = fluid.dygraph.to_variable(y_data)
label = fluid.layers.reshape(label, [args.batch_size, 1])
label.stop_gradient = True # 反向求导时，只有目标函数中的label部分对label求导（而不是label与img融合求导）
out, acc = train_model(img, label)
loss = fluid.layers.cross_entropy(out, label)
avg_loss = fluid.layers.mean(loss)
avg_loss.backward()
opt.minimize(avg_loss)
train_model.clear_gradients()
if batch_id % 100 == 0:
print("Loss at epoch {} step {}: {}, acc: {}".format(
i, batch_id, avg_loss.numpy(), acc.numpy()))
fluid.dygraph.save_dygraph(train_model.state_dict(),
args.save_dir + '/tsn_model')
print("Final loss: {}".format(avg_loss.numpy()))
def eval(args):
# parse config
config = parse_config(args.config)
val_config = merge_configs(config, 'valid', vars(args))
# print_configs(val_config, "Valid")
with fluid.dygraph.guard():
val_model = TSN(args.batch_size, 32)
label_dic = np.load('work/UCF-101_jpg/label_dir.npy',
allow_pickle=True).item()
label_dic = {v: k for k, v in label_dic.items()}
# get infer reader
val_reader = UCFReader('valid', val_config).create_reader()
# if no weight files specified, exit()
if args.weights:
weights = args.weights
else:
print("model path must be specified")
exit()
para_state_dict, _ = fluid.load_dygraph(weights)
val_model.load_dict(para_state_dict)
val_model.eval()
acc_list = []
for batch_id, data in enumerate(val_reader()):
dy_x_data = np.array([x[0] for x in data]).astype('float32')
y_data = np.array([[x[1]] for x in data]).astype('int64')
img = fluid.dygraph.to_variable(dy_x_data)
label = fluid.dygraph.to_variable(y_data)
label = fluid.layers.reshape(label, [args.batch_size, 1])
label.stop_gradient = True
out, acc = val_model(img, label)
print('batch_id=', batch_id, 'acc=', acc.numpy())
acc_list.append(acc.numpy())
print("验证集准确率为:{}".format(np.mean(acc_list)))
def infer(args):
# parse config
config = parse_config(args.config)
infer_config = merge_configs(config, 'infer', vars(args))
# print_configs(infer_config, "Infer")
with fluid.dygraph.guard():
infer_model = TSN(args.batch_size, 32)
label_dic = np.load('work/UCF-101_jpg/label_dir.npy', allow_pickle=True).item()
label_dic = {v: k for k, v in label_dic.items()}
# get infer reader
infer_reader = UCFReader('infer', infer_config).create_reader()
# if no weight files specified, exit()
if args.weights:
weights = args.weights
else:
print("model path must be specified")
exit()
para_state_dict, _ = fluid.load_dygraph(weights)
# print('para_state_dict:', para_state_dict)
infer_model.load_dict(para_state_dict)
infer_model.eval()
for batch_id, data in enumerate(infer_reader()):
dy_x_data = np.array([x[0] for x in data]).astype('float32')
y_data = np.array([[x[1]] for x in data])
img = fluid.dygraph.to_variable(dy_x_data)
out = infer_model(img)
label_id = fluid.layers.argmax(out, axis=1).numpy()[0]
print("实际标签{}, 预测结果{}".format(y_data, label_dic[label_id]))
