def forward_backward_v2(net, criterions, ctx, packet, is_train=True):
data, ht8, ht8_mask, paf8, paf8_mask = packet
criterion, criterion_ohkm = criterions
# split to gpus
data = gl.utils.split_and_load(data, ctx)
ht8 = gl.utils.split_and_load(ht8, ctx)
ht8_mask = gl.utils.split_and_load(ht8_mask, ctx)
paf8 = gl.utils.split_and_load(paf8, ctx)
paf8_mask = gl.utils.split_and_load(paf8_mask, ctx)
# run
ag.set_recording(is_train)
ag.set_training(is_train)
losses = []
for data_, ht8_, paf8_, ht8_mask_, paf8_mask_ in zip(
data, ht8, paf8, ht8_mask, paf8_mask):
# forward
out_ = net(data_)
losses_ = []
num_stage = len(out_)
for i in range(num_stage):
losses_.append(criterion(out_[i][0], ht8_, ht8_mask_))
losses_.append(criterion(out_[i][1], paf8_, paf8_mask_))
losses.append(losses_)
# backward
if is_train:
ag.backward(losses_)
ag.set_recording(False)
ag.set_training(False)
return losses
def export_model(self):
if not isinstance(self.model, nn.HybridBlock):
raise ValueError(
"Expected a HybridBlock but the model seems not one.")
autograd.set_training(False)
autograd.set_recording(False)
loader = self.create_dataloader("train")
raw_data = next(iter(loader))
splitted_data = utils.split_and_load(raw_data, self.ctx)
for data in splitted_data:
inputs, labels = self.parse_data(data, "train")
self.model(*inputs)
self.model.export(os.path.join(self.config.PARAM_DIR, "model"), 9999)
def forward_backward(net, criterions, ctx, data, rois, is_train=True):
criterion_cls1, criterion_cls2, criterion_reg = criterions
data = gl.utils.split_and_load(data, ctx)
rois = gl.utils.split_and_load(rois, ctx)
ag.set_recording(is_train)
ag.set_training(is_train)
# forward rpn
rpn_cls1, rpn_reg1, rpn_cls2, rpn_reg2 = [], [], [], []
for data_ in data:
rpn_cls1_, rpn_reg1_, rpn_cls2_, rpn_reg2_ = net(data_)
rpn_cls1.append(rpn_cls1_)
rpn_reg1.append(rpn_reg1_)
rpn_cls2.append(rpn_cls2_)
rpn_reg2.append(rpn_reg2_)
losses = []
anchor_proposals = net.anchor_proposals
for data_, rois_, rpn_cls1_, rpn_reg1_, rpn_cls2_, rpn_reg2_ in zip(
data, rois, rpn_cls1, rpn_reg1, rpn_cls2, rpn_reg2):
im_info = data_.shape[-2:]
# anchor target
# feat 1/8
# parallel stops here
batch_label1, batch_label_weight1, batch_bbox_targets1, batch_bbox_weights1 = anchor_proposals[
0].target(rpn_cls1_, rois_, im_info)
# loss cls
loss_cls1 = criterion_cls1(rpn_cls1_, batch_label1,
batch_label_weight1) / data_.shape[0]
# loss reg
loss_reg1 = criterion_reg(rpn_reg1_, batch_bbox_targets1,
batch_bbox_weights1) / data_.shape[0]
# feat 1/16
# parallel stops here
batch_label2, batch_label_weight2, batch_bbox_targets2, batch_bbox_weights2 = anchor_proposals[
1].target(rpn_cls2_, rois_, im_info)
# loss cls
loss_cls2 = criterion_cls2(rpn_cls2_, batch_label2,
batch_label_weight2) / data_.shape[0]
# loss reg
loss_reg2 = criterion_reg(rpn_reg2_, batch_bbox_targets2,
batch_bbox_weights2) / data_.shape[0]
loss = [loss_cls1, loss_reg1, loss_cls2, loss_reg2]
# backward
if is_train:
ag.backward(loss)
losses.append(loss)
ag.set_recording(False)
ag.set_training(False)
return losses
def forward_backward_v3(net, criterions, ctx, packet, is_train=True):
data, ht4, ht4_mask, paf4, paf4_mask, ht8, ht8_mask, paf8, paf8_mask, ht16, ht16_mask, paf16, paf16_mask = packet
criterion, criterion_ohkm = criterions
ht = [ht4, ht8, ht16]
paf = [paf4, paf8, paf16]
ht_mask = [ht4_mask, ht8_mask, ht16_mask]
paf_mask = [paf4_mask, paf8_mask, paf16_mask]
# split to gpus
data = gl.utils.split_and_load(data, ctx)
ht = [gl.utils.split_and_load(x, ctx) for x in ht]
paf = [gl.utils.split_and_load(x, ctx) for x in paf]
ht_mask = [gl.utils.split_and_load(x, ctx) for x in ht_mask]
paf_mask = [gl.utils.split_and_load(x, ctx) for x in paf_mask]
# run
ag.set_recording(is_train)
ag.set_training(is_train)
losses = []
for idx, data_ in enumerate(data):
# forward
g_ht4, g_paf4, r_ht4, r_paf4, g_ht8, g_paf8, r_ht8, r_paf8, g_ht16, g_paf16, r_ht16, r_paf16 = net(
data_)
ht4_, ht8_, ht16_ = [h[idx] for h in ht]
paf4_, paf8_, paf16_ = [p[idx] for p in paf]
ht4_mask_, ht8_mask_, ht16_mask_ = [hm[idx] for hm in ht_mask]
paf4_mask_, paf8_mask_, paf16_mask_ = [pm[idx] for pm in paf_mask]
# loss
losses_ = [
criterion(g_ht4, ht4_, ht4_mask_),
criterion_ohkm(r_ht4, ht4_, ht4_mask_),
criterion(g_ht8, ht8_, ht8_mask_),
criterion_ohkm(r_ht8, ht8_, ht8_mask_),
criterion(g_ht16, ht16_, ht16_mask_),
criterion_ohkm(r_ht16, ht16_, ht16_mask_),
criterion(g_paf4, paf4_, paf4_mask_),
criterion(r_paf4, paf4_, paf4_mask_),
criterion(g_paf8, paf8_, paf8_mask_),
criterion(r_paf8, paf8_, paf8_mask_),
criterion(g_paf16, paf16_, paf16_mask_),
criterion(r_paf16, paf16_, paf16_mask_)
]
losses.append(losses_)
# backward
if is_train:
ag.backward(losses_)
ag.set_recording(False)
ag.set_training(False)
return losses
def forward_backward(net, criterion, ctx, packet, is_train=True):
data, ht, mask = packet
data = gl.utils.split_and_load(data, ctx)
ht = gl.utils.split_and_load(ht, ctx)
mask = gl.utils.split_and_load(mask, ctx)
# run
ag.set_recording(is_train)
ag.set_training(is_train)
losses = []
for data_, ht_, mask_ in zip(data, ht, mask):
pred_ = net(data_)
losses_ = [criterion(ht_, pred_, mask_)]
losses.append(losses_)
if is_train:
ag.backward(losses_)
ag.set_recording(False)
ag.set_training(False)
return losses
def _process_epoch(self, mode):
color_code = esc_seq.GREEN if sys.platform != "win32" else ""
end_color_code = esc_seq.END if sys.platform != "win32" else ""
print(color_code +
"{}: epoch {:3d}/{:3d}".format(mode, self.latest_state +
1, self.config.MAX_EPOCHS) +
end_color_code)
loader = self.create_dataloader(mode)
handler = self.create_handler(mode=mode, num_batch=len(loader))
for i, raw_data in enumerate(loader):
gathered_outputs = []
gathered_losses = []
losses = []
tick = time.time()
splitted_data = utils.split_and_load(raw_data, self.ctx)
if mode == "train":
autograd.set_training(True)
autograd.set_recording(True)
elif mode == "test":
autograd.set_training(False)
autograd.set_recording(False)
for data in splitted_data:
inputs, labels = self.parse_data(data, mode)
outputs = self.parse_output(self.model(*inputs), mode)
gathered_outputs.append(outputs)
loss = self.compute_loss(outputs, labels)
gathered_losses.append(loss)
if mode == "train":
losses.extend(loss)
autograd.set_training(False)
autograd.set_recording(False)
if mode == "train":
autograd.backward(losses)
self.trainer.step(raw_data[0].shape[0])
handler.cleanup_batch(raw_data, gathered_outputs, gathered_losses,
i, tick)
handler.cleanup_epoch()