siamfc++中的trainer的具体实现
trainer = MODULES[name](optimizer, dataloader, monitors)
在videoanalyst/engine/trainer/trainer_base.py看看参数：
default_hyper_params = dict(
exp_name="default_training",
exp_save="snapshots",
max_epoch=20,
)
def __init__(self, optimizer, dataloader, monitors=[]):
self._hyper_params = deepcopy(
self.default_hyper_params) # mapping-like object
self._state = dict() # pipeline state
self._model = optimizer._model
self._losses = optimizer._model.loss
self._optimizer = optimizer
self._monitors = monitors
self._dataloader = iter(dataloader) # get the iterabel data loader
这里首先定义的monitors
monitors
videoanalyst/engine/monitor/monitor_impl/text_info.py直接print每个epoch中的学习率、loss、extra、time。
def update(self, engine_data: Dict):
r"""
"""
# state
engine_state = self._state["engine_state"]
# data
schedule_info = engine_data["schedule_info"]
training_losses = engine_data["training_losses"]
extras = engine_data["extras"]
time_dict = engine_data["time_dict"]
# schedule information
epoch = engine_state["epoch"]
print_str = 'epoch %d, ' % epoch
for k in schedule_info:
print_str += '%s: %.1e, ' % (k, schedule_info[k])
# loss info
for k in training_losses:
l = training_losses[k]
print_str += '%s: %.3f, ' % (k, l.detach().cpu().numpy())
# extra info
for extra in extras.values():
#if extra:
# extra = dist_utils.reduce_dict(extra)
for k in extra:
l = extra[k]
print_str += '%s: %.3f, ' % (k, l)
# pring elapsed time
for k in time_dict:
print_str += "%s: %.1e, " % (k, time_dict[k])
max_mem_mb = torch.cuda.max_memory_allocated() / 1024.0 / 1024.0
print_str += " max mem: {:.1f}M".format(max_mem_mb)
engine_state["print_str"] = print_str
同样videoanalyst/engine/monitor/monitor_impl/tensorboard_logger.py定义了update函数
def update(self, engine_data: Dict):
# from engine state calculate global step
engine_state = self._state["engine_state"]
epoch = engine_state["epoch"]
max_epoch = engine_state["max_epoch"]
iteration = engine_state["iteration"]
max_iteration = engine_state["max_iteration"]
global_step = iteration + epoch * max_iteration
# build at first update
if self._state["writer"] is None:
self._build_writer(global_step=global_step)
logger.info(
"Tensorboard writer built, starts recording from global_step=%d"
% global_step, )
logger.info(
"epoch=%d, max_epoch=%d, iteration=%d, max_iteration=%d" %
(epoch, max_epoch, iteration, max_iteration))
writer = self._state["writer"]
# traverse engine_data and put to scalar
self._add_scalar_recursively(writer, engine_data, "", global_step)
trainer
只需要看videoanalyst/engine/trainer/trainer_impl/regular_trainer.py
def train(self):
if not self._state["initialized"]:
self.init_train()
self._state["initialized"] = True
self._state["epoch"] += 1
epoch = self._state["epoch"]
num_iterations = self._hyper_params["num_iterations"]
# udpate engine_state
self._state["max_epoch"] = self._hyper_params["max_epoch"]
self._state["max_iteration"] = num_iterations
self._optimizer.modify_grad(epoch)
pbar = tqdm(range(num_iterations))
self._state["pbar"] = pbar
self._state["print_str"] = ""
time_dict = OrderedDict()
for iteration, _ in enumerate(pbar):
self._state["iteration"] = iteration
with Timer(name="data", output_dict=time_dict):
training_data = next(self._dataloader)
training_data = move_data_to_device(training_data,
self._state["devices"][0])
schedule_info = self._optimizer.schedule(epoch, iteration)
self._optimizer.zero_grad()
# forward propagation
with Timer(name="fwd", output_dict=time_dict):
predict_data = self._model(training_data)
training_losses, extras = OrderedDict(), OrderedDict()
for loss_name, loss in self._losses.items():
training_losses[loss_name], extras[loss_name] = loss(
predict_data, training_data)
total_loss = sum(training_losses.values())
# backward propagation
with Timer(name="bwd", output_dict=time_dict):
if self._optimizer.grad_scaler is not None:
self._optimizer.grad_scaler.scale(total_loss).backward()
else:
total_loss.backward()
self._optimizer.modify_grad(epoch, iteration)
with Timer(name="optim", output_dict=time_dict):
self._optimizer.step()
trainer_data = dict(
schedule_info=schedule_info,
training_losses=training_losses,
extras=extras,
time_dict=time_dict,
)
for monitor in self._monitors:
monitor.update(trainer_data)
del training_data
print_str = self._state["print_str"]
pbar.set_description(print_str)
已知trainingdata的形式是
training_data = dict(
im_z=im_z,
im_x=im_x,
bbox_z=bbox_z,
bbox_x=bbox_x,
cls_gt=cls_label,
ctr_gt=ctr_label,
box_gt=box_label,
is_negative_pair=int(is_negative_pair),
)
这里关注predict_data = self._model(training_data)是重点，在videoanalyst/model/task_model/taskmodel_impl/siamese_track.py定义了forward函数
def train_forward(self, training_data):
target_img = training_data["im_z"]
search_img = training_data["im_x"]
# backbone feature
f_z = self.basemodel(target_img)
f_x = self.basemodel(search_img)
# feature adjustment
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(
c_out, r_out)
predict_data = dict(
cls_pred=fcos_cls_score_final,
ctr_pred=fcos_ctr_score_final,
box_pred=fcos_bbox_final,
)
if self._hyper_params["corr_fea_output"]:
predict_data["corr_fea"] = corr_fea
return predict_data
根据pair图片对出结果predict_data。