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]))