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))
val_config = merge_configs(config, 'test', vars(args))
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
# 根据自己定义的网络，声明train_model
train_model = ResNet3D.generate_model(50)
if args.resume == True:
# 加载上一次训练的模型，继续训练
model, _ = fluid.dygraph.load_dygraph(args.save_dir + '/tsn_model')
train_model.load_dict(model)
print('Resume from ' + args.save_dir + '/tsn_model')
elif args.pretrain:
pretrain_weights = fluid.io.load_program_state(args.pretrain)
inner_state_dict = train_model.state_dict()
print('Resume from' + args.pretrain)
for name, para in inner_state_dict.items():
if ((name in pretrain_weights) and (not ('fc' in para.name))):
para.set_value(pretrain_weights[name])
else:
print('del ' + para.name)
opt = fluid.optimizer.Momentum(train_config.TRAIN.learning_rate,
train_config.TRAIN.learning_rate_decay,
parameter_list=train_model.parameters())
# build model
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# get reader
train_config.TRAIN.batch_size = train_config.TRAIN.batch_size
train_reader = KineticsReader(args.model_name.upper(), 'train',
train_config).create_reader()
val_reader = KineticsReader(args.model_name.upper(), 'valid',
val_config).create_reader()
epochs = args.epoch or train_model.epoch_num()
for i in range(epochs):
for batch_id, data in enumerate(train_reader(
)): # data (list) (batch)[seg_num,3 * seglen,size,size]
dy_x_data = np.array([x[0] for x in data]).astype(
'float32') # [batch, seg_num, 3 * seglen, size, size]
y_data = np.array([[x[1]] for x in data
]).astype('int64') # [batch, 1]
img = fluid.dygraph.to_variable(dy_x_data)
label = fluid.dygraph.to_variable(y_data)
label.stop_gradient = True
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()
logger.info("Loss at epoch {} step {}: {}, acc: {}".format(
i, batch_id, avg_loss.numpy(), acc.numpy()))
print("Loss at epoch {} step {}: {}, acc: {}".format(
i, batch_id, avg_loss.numpy(), acc.numpy()))
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.stop_gradient = True
out, acc = train_model(img, label)
acc_list.append(acc.numpy()[0])
logger.info("Val at epoch {}: acc: {}".format(
i, np.mean(acc_list)))
print("Val at epoch {}: acc: {}".format(i, np.mean(acc_list)) +
'\n')
if i % 10 == 0:
fluid.dygraph.save_dygraph(
train_model.state_dict(),
args.save_dir + '/tsn_model_' + str(i))
fluid.dygraph.save_dygraph(train_model.state_dict(),
args.save_dir + '/tsn_model')
logger.info("Final loss: {}".format(avg_loss.numpy()))
print("Final loss: {}".format(avg_loss.numpy()))
import torch
import paddle.fluid as fluid
from collections import OrderedDict
from model import ResNet3D
torch_weight = torch.load('r3d50_KM_200ep.pth',
map_location=torch.device('cpu'))
# 全连接1039
weight = []
for torch_key in torch_weight['state_dict'].keys():
weight.append(
[torch_key, torch_weight['state_dict'][torch_key].detach().numpy()])
with fluid.dygraph.guard():
# 加载网络结构
paddle_model = ResNet3D.generate_model(50)
# 读取paddle网络结构的参数列表
paddle_weight = paddle_model.state_dict()
# 将paddle权重的参数列表打印出来
# for paddle_key in paddle_weight:
# print(paddle_key)
# 进行模型参数转换
new_weight_dict = OrderedDict()
i = 0
for paddle_key in paddle_weight.keys():
print(paddle_key)
if 'num_batches_tracked' in weight[i][0]:
i += 1