def test_roi_align(self):
program = Program()
with program_guard(program):
x = layers.data(name="x", shape=[256, 30, 30], dtype="float32")
rois = layers.data(
name="rois", shape=[4], dtype="float32", lod_level=1)
output = layers.roi_align(x, rois, 14, 14, 0.5, 2)
self.assertIsNotNone(output)
print(str(program))
def forward(
self,
features,
rois,
):
# rois: list of ndarray (size: [n, 5]), p2 -> p5
# features: p5 -> p2
roi_features = []
roi_lvl_batch_count = []
for lvl in range(2, 6):
feature_map = features[5 - lvl]
spatial_scale = 1 / 2.**lvl
lvl_rois = rois[lvl - 2]
if len(lvl_rois) == 0:
continue
b = feature_map.shape[0]
batch_roi_count = []
for i in range(b):
batch_num = np.sum(lvl_rois[:, 0] == i)
batch_roi_count.append(batch_num)
roi_lvl_batch_count.append(batch_roi_count)
lvl_rois = F.create_lod_tensor(lvl_rois[:, 1:], [batch_roi_count],
place=F.CPUPlace())
lvl_rois = dg.to_variable(lvl_rois)
batch_roi_count = dg.to_variable(batch_roi_count).astype("int32")
lvl_roi_features = L.roi_align(
feature_map,
lvl_rois,
pooled_height=self.roi_size,
pooled_width=self.roi_size,
spatial_scale=spatial_scale,
sampling_ratio=self.roi_sampling_ratio,
rois_num=batch_roi_count)
c = lvl_roi_features.shape[1]
lvl_roi_features = L.reshape(
lvl_roi_features, (-1, self.roi_size * self.roi_size * c))
roi_features.append(lvl_roi_features)
roi_features = L.concat(roi_features, 0)
lvl_lod = [
0,
]
batch_lod = [
0,
]
for i in range(len(roi_lvl_batch_count)):
for j in range(len(roi_lvl_batch_count[i])):
batch_lod.append(batch_lod[-1] + roi_lvl_batch_count[i][j])
lvl_lod.append(lvl_lod[-1] + len(roi_lvl_batch_count[i]))
# roi_features.set_lod([lvl_lod, batch_lod])
fc6 = self.fc6(roi_features)
fc6 = L.relu(fc6)
fc7 = self.fc7(fc6)
fc7 = L.relu(fc7)
cls_score = self.cls_score(fc7)
cls_prob = L.softmax(cls_score)
bbox_pred = self.bbox_pred(fc7)
return cls_prob, bbox_pred, [lvl_lod, batch_lod]
def forward(self, features, rois):
# rois: list of ndarray (size: [n, 5]), p2 -> p5
# features: p5 -> p2
roi_features = []
roi_lvl_batch_count = []
for lvl in range(2, 6):
feature_map = features[5 - lvl]
spatial_scale = 1 / 2.**lvl
lvl_rois = rois[lvl - 2]
if len(lvl_rois) == 0:
continue
b = feature_map.shape[0]
batch_roi_count = []
for i in range(b):
batch_num = np.sum(lvl_rois[:, 0] == i)
batch_roi_count.append(batch_num)
roi_lvl_batch_count.append(batch_roi_count)
lvl_rois = F.create_lod_tensor(lvl_rois[:, 1:].astype("float32"),
[batch_roi_count],
place=F.CPUPlace())
lvl_rois = dg.to_variable(lvl_rois)
batch_roi_count = dg.to_variable(batch_roi_count).astype("int32")
lvl_roi_features = L.roi_align(
feature_map,
lvl_rois,
pooled_height=self.roi_size,
pooled_width=self.roi_size,
spatial_scale=spatial_scale,
sampling_ratio=self.roi_sampling_ratio,
rois_num=batch_roi_count)
roi_features.append(lvl_roi_features)
if roi_features == []:
return None, None, None, None, None
roi_features = L.concat(roi_features) # (n, c, 14, 14)
for conv in self.convs:
roi_features = conv(roi_features)
roi_features = L.relu(roi_features)
ann_index_lowres = self.AnnIndex_lowres(roi_features)
index_UV_lowres = self.Index_UV_lowres(roi_features)
u_lowres = self.U_lowres(roi_features)
v_lowres = self.V_lowres(roi_features)
b, _, h, w = ann_index_lowres.shape
zero_out = dg.to_variable(np.zeros((b, 10, h, w)).astype("float32"))
ann_index_lowres = L.concat([ann_index_lowres, zero_out], 1)
ann_index = self.AnnIndex(ann_index_lowres)
index = self.Index_UV(index_UV_lowres)
u = self.U_estimated(u_lowres)
v = self.V_estimated(v_lowres)
lvl_lod = [
0,
]
batch_lod = [
0,
]
for i in range(len(roi_lvl_batch_count)):
for j in range(len(roi_lvl_batch_count[i])):
batch_lod.append(batch_lod[-1] + roi_lvl_batch_count[i][j])
lvl_lod.append(lvl_lod[-1] + len(roi_lvl_batch_count[i]))
# roi_features.set_lod([lvl_lod, batch_lod])
return ann_index, index, u, v, [lvl_lod, batch_lod]