HMN项目代码
本周主要对HMN项目中的:
main.py,train.py,eval.py,data_loader.py,build_loaders.py
进行了学习,代码的学习情况如下。
整个项目的main函数,包含了项目的整体流程。
if __name__ == '__main__':
# 初始化文件保存的内容信息、路径信息
info, path_join = init_log()
cfgs = get_settings()
set_random_seed(seed=cfgs.seed)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_loader, valid_loader, test_loader = build_loaders(cfgs)
hungary_matcher = HungarianMatcher() # 匈牙利算法
model = build_model(cfgs) # 搭建模型
model = model.float() # 转换模型的所有parameters和buffers的类型
model.to(device) # 将模型加载到 cuda 上
model = train_fn(cfgs, cfgs.model_name, model, hungary_matcher, train_loader, valid_loader, device, info, path_join) # 训练模型
model.load_state_dict(torch.load(cfgs.train.save_checkpoints_path))
model.eval()
整体流程为:
对于网络各个参数的设置,见文件setting.py,其对部分参数也进行了注释介绍,在这里面有几个比较重要的参数需要注意:
batchsize默认值为64,drop_prob(drop比例)默认值为0.5,sample_numb(给定一个视频抽取帧的数目)默认值为15。
加载完各个参数的值,接下来是加载数据:
分别定义了加载训练、验证、测试数据的函数,其中DataLoader方法为pytorch自带的方法,CaptionDataset类在data_loader.py文件中定义。
def get_test_loader(cfgs: TotalConfigs, is_total=False):
test_dataset = CaptionDataset(cfgs=cfgs, mode='test', save_on_disk=True, is_total=is_total)
test_loader = DataLoader(dataset=test_dataset, batch_size=cfgs.bsz,
shuffle=False, collate_fn=collate_fn_caption,
num_workers=0)
return test_loader
def get_train_loader(cfgs: TotalConfigs, save_on_disk=True):
train_dataset = CaptionDataset(cfgs=cfgs, mode='train', save_on_disk=save_on_disk, is_total=False)
train_loader = DataLoader(dataset=train_dataset, batch_size=cfgs.bsz,
shuffle=False, collate_fn=collate_fn_caption,
num_workers=0)
return train_loader
def get_valid_loader(cfgs: TotalConfigs, is_total=False):
valid_dataset = CaptionDataset(cfgs=cfgs, mode='valid', save_on_disk=True, is_total=is_total)
valid_loader = DataLoader(dataset=valid_dataset, batch_size=cfgs.bsz,
shuffle=False, collate_fn=collate_fn_caption,
num_workers=0)
return valid_loader
在该文件中,创造了CaptionDataset类,定义了模型加载language信息和visual信息的属性和方法,通过h5py文件加载视频的2d、3d特征,最后再获取视频的字幕、语义信息。
在经过 1、2 两步操作后,已经确定好了网络模型的一些参数信息,并且封装好了数据,接下来就是搭建模型。
可以看到模型由编码器和解码器构成,其中编码器包含了transformer、entity_level_encoder、predicate_level_encoder、sentence_level_encoder
四个模块,解码器则只通过一个decoder实现。
if model_name == 'HMN':
# encoders
transformer = Transformer(d_model=d_model, nhead=nheads,
num_encoder_layers=trans_num_encoder_layers,
num_decoder_layers=trans_num_decoder_layers,
dim_feedforward=dim_feedforward,
dropout=trans_dropout,
activation=transformer_activation)
entity_level_encoder = EntityLevelEncoder(transformer=transformer,
max_objects=max_objects,
object_dim=object_dim,
feature2d_dim=feature2d_dim,
feature3d_dim=feature3d_dim,
hidden_dim=hidden_dim,
word_dim=semantics_dim)
predicate_level_encoder = PredicateLevelEncoder(feature3d_dim=feature3d_dim,
hidden_dim=hidden_dim,
semantics_dim=semantics_dim,
useless_objects=False)
sentence_level_encoder = SentenceLevelEncoder(feature2d_dim=feature2d_dim,
hidden_dim=hidden_dim,
semantics_dim=semantics_dim,
useless_objects=False)
# decoder
decoder = Decoder(semantics_dim=semantics_dim, hidden_dim=hidden_dim,
num_layers=decoder_num_layers, embed_dim=embed_dim, n_vocab=n_vocab,
with_objects=True, with_action=True, with_video=True,
with_objects_semantics=True,
with_action_semantics=True,
with_video_semantics=True)
在settings.py文件中,上述参数的设定值如下:
parser.add_argument('--nheads', type=int, default=8)
parser.add_argument('--entity_encoder_layer', type=int, default=2)
parser.add_argument('--entity_decoder_layer', type=int, default=2)
parser.add_argument('--dim_feedforward', type=int, default=2048)
parser.add_argument('--transformer_activation', type=str, default='relu')
parser.add_argument('--d_model', type=int, default=512)
parser.add_argument('--transformer_dropout', type=float, default=0.1)
搭建好模型后便可以对模型进行训练,基于已装载的数据和搭建好的网络,整个模型的训练过程如下(伪代码):
def train_fn(cfgs: TotalConfigs, model_name: str, model: nn.Module, matcher: HungarianMatcher, train_loader,
valid_loader, device, info, path_join):
# 加载 loss 函数
loss = set_loss_funcand_state_list()
# 加载二进制文件
pickle.load(f)
print('===================Training begin====================')
for epoch in range(max_epoch):
for i,item in enumerate(train_loader):
"""
1.加载缓冲区数据
2.创建优化器
3.进行前向传播
4.计算当前loss
5.进行方向传播
6.保存loss到列表中
7.梯度裁剪
8.优化器更新
9.内循环每循环X次保存一次checkpoint
"""
lr_scheduler.step() # 外循环更新学习率参数
print('===================Training is finished====================')
return model
在保存checkpoint文件时,需要使用验证集对当前模型进行分数评估,评估函数在eavl.py中进行定义。
在该文件中,定义了eavl_fn函数来返回分数评估结果,定义了language_eval函数来根据predictions和groundtruth计算得分。
def eval_fn() -> dict:
model.eval() #关闭Batch Normalization 和 Dropout,保证评估过程中BN层的均值和方差不变
for i item in enumerate(loader):
"""
1.加载2d、3d特征;
2.使用模型对结果进行预测
3.将预测结果和对应的groundtruth保存到列表中
"""
model.train() #启用batch normalization 和 dropout,进行下一轮训练更新参数
# 计算预测结果与groundtruth对比的得分
score_states = language_eval(sample_seqs=predictions, groundtruth_seqs=gts, vids_list=vids_list)
return score_states
该项目总共训练20个epoch,在每一个epoch的训练过程中,train文件中的内循环每循环500次,就会计算一次得分(bleu4,cider,meteor,rouge),如果当前的cider分数为最高则更新一次训练权重。
def save_checkpoint(best_score, cfgs, cnt, device, epoch, i, idx2word, info, model, valid_loader, vid2groundtruth,
is_last_epoch, path_join):
if cnt % cfgs.train.save_checkpoints_every == 0: # enumerate_for循环,每循环500次 save_checkpoint
ckpt_path = cfgs.train.save_checkpoints_path
# eavl.py 文件计算分数
scores = eval_fn(model=model, loader=valid_loader, device=device,
idx2word=idx2word, save_on_disk=False, cfgs=cfgs,
vid2groundtruth=vid2groundtruth, is_last_epoch=is_last_epoch, path_join=path_join)
cider_score = scores['CIDEr']
if best_score is None or cider_score > best_score:
best_score = cider_score
torch.save(model.state_dict(), ckpt_path)
info('=' * 10 + '[EPOCH{epoch} iter{it}] :Best Cider is {bs}, Current Cider is {cs}'.
format(epoch=epoch, it=i, bs=best_score, cs=cider_score) + '=' * 10)