test_bmn.py 30.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
16 17
import os
import tempfile
18
import unittest
19 20 21 22

import numpy as np
from predictor_utils import PredictorTools

L
Leo Chen 已提交
23
import paddle
24 25 26
import paddle.fluid as fluid
from paddle.fluid import ParamAttr
from paddle.fluid.dygraph import to_variable
27
from paddle.fluid.dygraph.io import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX
28
from paddle.jit import ProgramTranslator, to_static
29

M
MRXLT 已提交
30
SEED = 2000
31 32 33 34 35 36 37 38 39 40
DATATYPE = 'float32'
program_translator = ProgramTranslator()

# Note: Set True to eliminate randomness.
#     1. For one operation, cuDNN has several algorithms,
#        some algorithm results are non-deterministic, like convolution algorithms.
if fluid.is_compiled_with_cuda():
    fluid.set_flags({'FLAGS_cudnn_deterministic': True})


41 42 43 44
def get_interp1d_mask(
    tscale, dscale, prop_boundary_ratio, num_sample, num_sample_perbin
):
    """generate sample mask for each point in Boundary-Matching Map"""
45 46 47 48 49 50 51 52 53 54
    mask_mat = []
    for start_index in range(tscale):
        mask_mat_vector = []
        for duration_index in range(dscale):
            if start_index + duration_index < tscale:
                p_xmin = start_index
                p_xmax = start_index + duration_index
                center_len = float(p_xmax - p_xmin) + 1
                sample_xmin = p_xmin - center_len * prop_boundary_ratio
                sample_xmax = p_xmax + center_len * prop_boundary_ratio
55 56 57 58 59 60 61
                p_mask = _get_interp1d_bin_mask(
                    sample_xmin,
                    sample_xmax,
                    tscale,
                    num_sample,
                    num_sample_perbin,
                )
62 63 64 65 66 67 68 69 70 71 72 73
            else:
                p_mask = np.zeros([tscale, num_sample])
            mask_mat_vector.append(p_mask)
        mask_mat_vector = np.stack(mask_mat_vector, axis=2)
        mask_mat.append(mask_mat_vector)
    mask_mat = np.stack(mask_mat, axis=3)
    mask_mat = mask_mat.astype(np.float32)

    sample_mask = np.reshape(mask_mat, [tscale, -1])
    return sample_mask


74 75 76 77
def _get_interp1d_bin_mask(
    seg_xmin, seg_xmax, tscale, num_sample, num_sample_perbin
):
    """generate sample mask for a boundary-matching pair"""
78 79 80 81 82 83 84 85
    plen = float(seg_xmax - seg_xmin)
    plen_sample = plen / (num_sample * num_sample_perbin - 1.0)
    total_samples = [
        seg_xmin + plen_sample * ii
        for ii in range(num_sample * num_sample_perbin)
    ]
    p_mask = []
    for idx in range(num_sample):
86 87 88
        bin_samples = total_samples[
            idx * num_sample_perbin : (idx + 1) * num_sample_perbin
        ]
89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
        bin_vector = np.zeros([tscale])
        for sample in bin_samples:
            sample_upper = math.ceil(sample)
            sample_decimal, sample_down = math.modf(sample)
            if int(sample_down) <= (tscale - 1) and int(sample_down) >= 0:
                bin_vector[int(sample_down)] += 1 - sample_decimal
            if int(sample_upper) <= (tscale - 1) and int(sample_upper) >= 0:
                bin_vector[int(sample_upper)] += sample_decimal
        bin_vector = 1.0 / num_sample_perbin * bin_vector
        p_mask.append(bin_vector)
    p_mask = np.stack(p_mask, axis=1)
    return p_mask


class Conv1D(fluid.dygraph.Layer):
104 105 106 107 108 109 110 111 112 113
    def __init__(
        self,
        prefix,
        num_channels=256,
        num_filters=256,
        size_k=3,
        padding=1,
        groups=1,
        act="relu",
    ):
114
        super().__init__()
115
        fan_in = num_channels * size_k * 1
116 117 118 119 120 121 122 123 124 125
        k = 1.0 / math.sqrt(fan_in)
        param_attr = ParamAttr(
            name=prefix + "_w",
            initializer=fluid.initializer.Uniform(low=-k, high=k),
        )
        bias_attr = ParamAttr(
            name=prefix + "_b",
            initializer=fluid.initializer.Uniform(low=-k, high=k),
        )

126 127 128 129
        self._conv2d = paddle.nn.Conv2D(
            in_channels=num_channels,
            out_channels=num_filters,
            kernel_size=(1, size_k),
130 131 132
            stride=1,
            padding=(0, padding),
            groups=groups,
133
            weight_attr=param_attr,
134 135
            bias_attr=bias_attr,
        )
136 137 138 139

    def forward(self, x):
        x = fluid.layers.unsqueeze(input=x, axes=[2])
        x = self._conv2d(x)
140
        x = paddle.squeeze(x, axis=[2])
141 142 143 144 145
        return x


class BMN(fluid.dygraph.Layer):
    def __init__(self, cfg):
146
        super().__init__()
147 148 149 150 151 152 153 154 155 156 157 158

        self.tscale = cfg.tscale
        self.dscale = cfg.dscale
        self.prop_boundary_ratio = cfg.prop_boundary_ratio
        self.num_sample = cfg.num_sample
        self.num_sample_perbin = cfg.num_sample_perbin

        self.hidden_dim_1d = 256
        self.hidden_dim_2d = 128
        self.hidden_dim_3d = 512

        # Base Module
159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
        self.b_conv1 = Conv1D(
            prefix="Base_1",
            num_channels=cfg.feat_dim,
            num_filters=self.hidden_dim_1d,
            size_k=3,
            padding=1,
            groups=4,
            act="relu",
        )
        self.b_conv2 = Conv1D(
            prefix="Base_2",
            num_filters=self.hidden_dim_1d,
            size_k=3,
            padding=1,
            groups=4,
            act="relu",
        )
176 177

        # Temporal Evaluation Module
178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
        self.ts_conv1 = Conv1D(
            prefix="TEM_s1",
            num_filters=self.hidden_dim_1d,
            size_k=3,
            padding=1,
            groups=4,
            act="relu",
        )
        self.ts_conv2 = Conv1D(
            prefix="TEM_s2", num_filters=1, size_k=1, padding=0, act="sigmoid"
        )
        self.te_conv1 = Conv1D(
            prefix="TEM_e1",
            num_filters=self.hidden_dim_1d,
            size_k=3,
            padding=1,
            groups=4,
            act="relu",
        )
        self.te_conv2 = Conv1D(
            prefix="TEM_e2", num_filters=1, size_k=1, padding=0, act="sigmoid"
        )

        # Proposal Evaluation Module
        self.p_conv1 = Conv1D(
            prefix="PEM_1d",
            num_filters=self.hidden_dim_2d,
            size_k=3,
            padding=1,
            act="relu",
        )
209 210

        # init to speed up
211 212 213 214 215 216 217
        sample_mask = get_interp1d_mask(
            self.tscale,
            self.dscale,
            self.prop_boundary_ratio,
            self.num_sample,
            self.num_sample_perbin,
        )
218 219
        self.sample_mask = fluid.dygraph.base.to_variable(sample_mask)
        self.sample_mask.stop_gradient = True
220

221 222 223 224
        self.p_conv3d1 = paddle.nn.Conv3D(
            in_channels=128,
            out_channels=self.hidden_dim_3d,
            kernel_size=(self.num_sample, 1, 1),
225 226
            stride=(self.num_sample, 1, 1),
            padding=0,
227 228
            weight_attr=paddle.ParamAttr(name="PEM_3d1_w"),
            bias_attr=paddle.ParamAttr(name="PEM_3d1_b"),
229
        )
230

231 232 233 234
        self.p_conv2d1 = paddle.nn.Conv2D(
            in_channels=512,
            out_channels=self.hidden_dim_2d,
            kernel_size=1,
235 236
            stride=1,
            padding=0,
237
            weight_attr=ParamAttr(name="PEM_2d1_w"),
238 239
            bias_attr=ParamAttr(name="PEM_2d1_b"),
        )
240 241 242 243
        self.p_conv2d2 = paddle.nn.Conv2D(
            in_channels=128,
            out_channels=self.hidden_dim_2d,
            kernel_size=3,
244 245
            stride=1,
            padding=1,
246
            weight_attr=ParamAttr(name="PEM_2d2_w"),
247 248
            bias_attr=ParamAttr(name="PEM_2d2_b"),
        )
249 250 251 252
        self.p_conv2d3 = paddle.nn.Conv2D(
            in_channels=128,
            out_channels=self.hidden_dim_2d,
            kernel_size=3,
253 254
            stride=1,
            padding=1,
255
            weight_attr=ParamAttr(name="PEM_2d3_w"),
256 257
            bias_attr=ParamAttr(name="PEM_2d3_b"),
        )
258 259 260 261
        self.p_conv2d4 = paddle.nn.Conv2D(
            in_channels=128,
            out_channels=2,
            kernel_size=1,
262 263
            stride=1,
            padding=0,
264
            weight_attr=ParamAttr(name="PEM_2d4_w"),
265 266
            bias_attr=ParamAttr(name="PEM_2d4_b"),
        )
267

A
Aurelius84 已提交
268
    @to_static
269 270
    def forward(self, x):
        # Base Module
271 272
        x = paddle.nn.functional.relu(self.b_conv1(x))
        x = paddle.nn.functional.relu(self.b_conv2(x))
273 274

        # TEM
275 276
        xs = paddle.nn.functional.relu(self.ts_conv1(x))
        xs = paddle.nn.functional.relu(self.ts_conv2(xs))
277
        xs = paddle.squeeze(xs, axis=[1])
278 279
        xe = paddle.nn.functional.relu(self.te_conv1(x))
        xe = paddle.nn.functional.relu(self.te_conv2(xe))
280
        xe = paddle.squeeze(xe, axis=[1])
281 282

        # PEM
283
        xp = paddle.nn.functional.relu(self.p_conv1(x))
284
        # BM layer
285
        xp = fluid.layers.matmul(xp, self.sample_mask)
286
        xp = paddle.reshape(xp, shape=[0, 0, -1, self.dscale, self.tscale])
287 288

        xp = self.p_conv3d1(xp)
289
        xp = paddle.tanh(xp)
290
        xp = paddle.squeeze(xp, axis=[2])
291 292 293 294
        xp = paddle.nn.functional.relu(self.p_conv2d1(xp))
        xp = paddle.nn.functional.relu(self.p_conv2d2(xp))
        xp = paddle.nn.functional.relu(self.p_conv2d3(xp))
        xp = paddle.nn.functional.sigmoid(self.p_conv2d4(xp))
295 296 297
        return xp, xs, xe


298 299 300
def bmn_loss_func(
    pred_bm, pred_start, pred_end, gt_iou_map, gt_start, gt_end, cfg
):
301 302 303 304 305
    def _get_mask(cfg):
        dscale = cfg.dscale
        tscale = cfg.tscale
        bm_mask = []
        for idx in range(dscale):
306 307 308
            mask_vector = [1 for i in range(tscale - idx)] + [
                0 for i in range(idx)
            ]
309 310
            bm_mask.append(mask_vector)
        bm_mask = np.array(bm_mask, dtype=np.float32)
311 312 313
        self_bm_mask = fluid.layers.create_global_var(
            shape=[dscale, tscale], value=0, dtype=DATATYPE, persistable=True
        )
314 315 316 317 318 319
        fluid.layers.assign(bm_mask, self_bm_mask)
        self_bm_mask.stop_gradient = True
        return self_bm_mask

    def tem_loss_func(pred_start, pred_end, gt_start, gt_end):
        def bi_loss(pred_score, gt_label):
320 321
            pred_score = paddle.reshape(x=pred_score, shape=[-1])
            gt_label = paddle.reshape(x=gt_label, shape=[-1])
322 323
            gt_label.stop_gradient = True
            pmask = fluid.layers.cast(x=(gt_label > 0.5), dtype=DATATYPE)
324 325 326
            num_entries = fluid.layers.cast(
                fluid.layers.shape(pmask), dtype=DATATYPE
            )
327
            num_positive = fluid.layers.cast(paddle.sum(pmask), dtype=DATATYPE)
328 329 330 331 332
            ratio = num_entries / num_positive
            coef_0 = 0.5 * ratio / (ratio - 1)
            coef_1 = 0.5 * ratio
            epsilon = 0.000001
            # temp = fluid.layers.log(pred_score + epsilon)
333
            loss_pos = paddle.multiply(
334 335
                fluid.layers.log(pred_score + epsilon), pmask
            )
336
            loss_pos = coef_1 * fluid.layers.reduce_mean(loss_pos)
337
            loss_neg = paddle.multiply(
338 339
                fluid.layers.log(1.0 - pred_score + epsilon), (1.0 - pmask)
            )
340 341 342 343 344 345 346 347 348 349 350
            loss_neg = coef_0 * fluid.layers.reduce_mean(loss_neg)
            loss = -1 * (loss_pos + loss_neg)
            return loss

        loss_start = bi_loss(pred_start, gt_start)
        loss_end = bi_loss(pred_end, gt_end)
        loss = loss_start + loss_end
        return loss

    def pem_reg_loss_func(pred_score, gt_iou_map, mask):

351
        gt_iou_map = paddle.multiply(gt_iou_map, mask)
352 353

        u_hmask = fluid.layers.cast(x=gt_iou_map > 0.7, dtype=DATATYPE)
354
        u_mmask = paddle.logical_and(gt_iou_map <= 0.7, gt_iou_map > 0.3)
355
        u_mmask = fluid.layers.cast(x=u_mmask, dtype=DATATYPE)
356
        u_lmask = paddle.logical_and(gt_iou_map <= 0.3, gt_iou_map >= 0.0)
357
        u_lmask = fluid.layers.cast(x=u_lmask, dtype=DATATYPE)
358
        u_lmask = paddle.multiply(u_lmask, mask)
359

360 361 362
        num_h = fluid.layers.cast(paddle.sum(u_hmask), dtype=DATATYPE)
        num_m = fluid.layers.cast(paddle.sum(u_mmask), dtype=DATATYPE)
        num_l = fluid.layers.cast(paddle.sum(u_lmask), dtype=DATATYPE)
363 364 365

        r_m = num_h / num_m
        u_smmask = fluid.layers.assign(
366
            local_random.uniform(
367 368 369
                0.0, 1.0, [gt_iou_map.shape[1], gt_iou_map.shape[2]]
            ).astype(DATATYPE)
        )
370
        u_smmask = paddle.multiply(u_mmask, u_smmask)
371
        u_smmask = fluid.layers.cast(x=(u_smmask > (1.0 - r_m)), dtype=DATATYPE)
372 373 374

        r_l = num_h / num_l
        u_slmask = fluid.layers.assign(
375
            local_random.uniform(
376 377 378
                0.0, 1.0, [gt_iou_map.shape[1], gt_iou_map.shape[2]]
            ).astype(DATATYPE)
        )
379
        u_slmask = paddle.multiply(u_lmask, u_slmask)
380
        u_slmask = fluid.layers.cast(x=(u_slmask > (1.0 - r_l)), dtype=DATATYPE)
381 382 383 384

        weights = u_hmask + u_smmask + u_slmask
        weights.stop_gradient = True
        loss = fluid.layers.square_error_cost(pred_score, gt_iou_map)
385
        loss = paddle.multiply(loss, weights)
386
        loss = 0.5 * paddle.sum(loss) / paddle.sum(weights)
387 388 389 390

        return loss

    def pem_cls_loss_func(pred_score, gt_iou_map, mask):
391
        gt_iou_map = paddle.multiply(gt_iou_map, mask)
392 393 394
        gt_iou_map.stop_gradient = True
        pmask = fluid.layers.cast(x=(gt_iou_map > 0.9), dtype=DATATYPE)
        nmask = fluid.layers.cast(x=(gt_iou_map <= 0.9), dtype=DATATYPE)
395
        nmask = paddle.multiply(nmask, mask)
396

397 398
        num_positive = paddle.sum(pmask)
        num_entries = num_positive + paddle.sum(nmask)
399 400 401 402
        ratio = num_entries / num_positive
        coef_0 = 0.5 * ratio / (ratio - 1)
        coef_1 = 0.5 * ratio
        epsilon = 0.000001
403
        loss_pos = paddle.multiply(
404 405
            fluid.layers.log(pred_score + epsilon), pmask
        )
406
        loss_pos = coef_1 * paddle.sum(loss_pos)
407
        loss_neg = paddle.multiply(
408 409
            fluid.layers.log(1.0 - pred_score + epsilon), nmask
        )
410
        loss_neg = coef_0 * paddle.sum(loss_neg)
411 412 413
        loss = -1 * (loss_pos + loss_neg) / num_entries
        return loss

414
    pred_bm_reg = paddle.squeeze(
2
201716010711 已提交
415
        paddle.slice(pred_bm, axes=[1], starts=[0], ends=[1]), axis=[1]
416
    )
417
    pred_bm_cls = paddle.squeeze(
2
201716010711 已提交
418
        paddle.slice(pred_bm, axes=[1], starts=[1], ends=[2]), axis=[1]
419
    )
420 421 422 423 424 425 426 427 428 429 430 431

    bm_mask = _get_mask(cfg)

    pem_reg_loss = pem_reg_loss_func(pred_bm_reg, gt_iou_map, bm_mask)
    pem_cls_loss = pem_cls_loss_func(pred_bm_cls, gt_iou_map, bm_mask)

    tem_loss = tem_loss_func(pred_start, pred_end, gt_start, gt_end)

    loss = tem_loss + 10 * pem_reg_loss + pem_cls_loss
    return loss, tem_loss, pem_reg_loss, pem_cls_loss


432
class Args:
433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458
    epoch = 1
    batch_size = 4
    learning_rate = 0.1
    learning_rate_decay = 0.1
    lr_decay_iter = 4200
    l2_weight_decay = 1e-4
    valid_interval = 20
    log_interval = 5
    train_batch_num = valid_interval
    valid_batch_num = 5

    tscale = 50
    dscale = 50
    feat_dim = 100
    prop_boundary_ratio = 0.5
    num_sample = 2
    num_sample_perbin = 2


def optimizer(cfg, parameter_list):
    bd = [cfg.lr_decay_iter]
    base_lr = cfg.learning_rate
    lr_decay = cfg.learning_rate_decay
    l2_weight_decay = cfg.l2_weight_decay
    lr = [base_lr, base_lr * lr_decay]
    optimizer = fluid.optimizer.Adam(
459
        fluid.layers.piecewise_decay(boundaries=bd, values=lr),
460 461
        parameter_list=parameter_list,
        regularization=fluid.regularizer.L2DecayRegularizer(
462 463 464
            regularization_coeff=l2_weight_decay
        ),
    )
465 466 467 468 469
    return optimizer


def fake_data_reader(args, mode='train'):
    def iou_with_anchors(anchors_min, anchors_max, box_min, box_max):
470
        """Compute jaccard score between a box and the anchors."""
471 472 473
        len_anchors = anchors_max - anchors_min
        int_xmin = np.maximum(anchors_min, box_min)
        int_xmax = np.minimum(anchors_max, box_max)
474
        inter_len = np.maximum(int_xmax - int_xmin, 0.0)
475 476 477 478 479
        union_len = len_anchors - inter_len + box_max - box_min
        jaccard = np.divide(inter_len, union_len)
        return jaccard

    def ioa_with_anchors(anchors_min, anchors_max, box_min, box_max):
480
        """Compute intersection between score a box and the anchors."""
481 482 483
        len_anchors = anchors_max - anchors_min
        int_xmin = np.maximum(anchors_min, box_min)
        int_xmax = np.minimum(anchors_max, box_max)
484
        inter_len = np.maximum(int_xmax - int_xmin, 0.0)
485 486 487 488 489
        scores = np.divide(inter_len, len_anchors)
        return scores

    def get_match_map(tscale):
        match_map = []
490
        tgap = 1.0 / tscale
491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513
        for idx in range(tscale):
            tmp_match_window = []
            xmin = tgap * idx
            for jdx in range(1, tscale + 1):
                xmax = xmin + tgap * jdx
                tmp_match_window.append([xmin, xmax])
            match_map.append(tmp_match_window)
        match_map = np.array(match_map)
        match_map = np.transpose(match_map, [1, 0, 2])
        match_map = np.reshape(match_map, [-1, 2])
        match_map = match_map
        anchor_xmin = [tgap * i for i in range(tscale)]
        anchor_xmax = [tgap * i for i in range(1, tscale + 1)]

        return match_map, anchor_xmin, anchor_xmax

    def get_video_label(match_map, anchor_xmin, anchor_xmax):
        video_second = local_random.randint(75, 90)
        label_num = local_random.randint(1, 3)

        gt_bbox = []
        gt_iou_map = []
        for idx in range(label_num):
514 515 516 517 518 519
            duration = local_random.uniform(
                video_second * 0.4, video_second * 0.8
            )
            start_t = local_random.uniform(
                0.1 * video_second, video_second - duration
            )
520 521 522
            tmp_start = max(min(1, start_t / video_second), 0)
            tmp_end = max(min(1, (start_t + duration) / video_second), 0)
            gt_bbox.append([tmp_start, tmp_end])
523 524 525 526 527 528
            tmp_gt_iou_map = iou_with_anchors(
                match_map[:, 0], match_map[:, 1], tmp_start, tmp_end
            )
            tmp_gt_iou_map = np.reshape(
                tmp_gt_iou_map, [args.dscale, args.tscale]
            )
529 530 531 532 533 534 535
            gt_iou_map.append(tmp_gt_iou_map)
        gt_iou_map = np.array(gt_iou_map)
        gt_iou_map = np.max(gt_iou_map, axis=0)

        gt_bbox = np.array(gt_bbox)
        gt_xmins = gt_bbox[:, 0]
        gt_xmaxs = gt_bbox[:, 1]
536
        gt_len_small = 3.0 / args.tscale
537
        gt_start_bboxs = np.stack(
538 539
            (gt_xmins - gt_len_small / 2, gt_xmins + gt_len_small / 2), axis=1
        )
540
        gt_end_bboxs = np.stack(
541 542
            (gt_xmaxs - gt_len_small / 2, gt_xmaxs + gt_len_small / 2), axis=1
        )
543 544 545 546 547

        match_score_start = []
        for jdx in range(len(anchor_xmin)):
            match_score_start.append(
                np.max(
548 549 550 551 552 553 554 555
                    ioa_with_anchors(
                        anchor_xmin[jdx],
                        anchor_xmax[jdx],
                        gt_start_bboxs[:, 0],
                        gt_start_bboxs[:, 1],
                    )
                )
            )
556 557 558 559
        match_score_end = []
        for jdx in range(len(anchor_xmin)):
            match_score_end.append(
                np.max(
560 561 562 563 564 565 566 567
                    ioa_with_anchors(
                        anchor_xmin[jdx],
                        anchor_xmax[jdx],
                        gt_end_bboxs[:, 0],
                        gt_end_bboxs[:, 1],
                    )
                )
            )
568 569 570 571 572 573 574 575 576 577 578 579

        gt_start = np.array(match_score_start)
        gt_end = np.array(match_score_end)
        return gt_iou_map, gt_start, gt_end

    def reader():
        batch_out = []
        iter_num = args.batch_size * 100
        match_map, anchor_xmin, anchor_xmax = get_match_map(args.tscale)

        for video_idx in range(iter_num):
            video_feat = local_random.random_sample(
580 581
                [args.feat_dim, args.tscale]
            ).astype('float32')
582
            gt_iou_map, gt_start, gt_end = get_video_label(
583 584
                match_map, anchor_xmin, anchor_xmax
            )
585 586 587 588 589

            if mode == 'train' or mode == 'valid':
                batch_out.append((video_feat, gt_iou_map, gt_start, gt_end))
            elif mode == 'test':
                batch_out.append(
590 591
                    (video_feat, gt_iou_map, gt_start, gt_end, video_idx)
                )
592
            else:
593
                raise NotImplementedError(
594 595
                    'mode {} not implemented'.format(mode)
                )
596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624
            if len(batch_out) == args.batch_size:
                yield batch_out
                batch_out = []

    return reader


# Validation
def val_bmn(model, args):
    val_reader = fake_data_reader(args, 'valid')

    loss_data = []
    for batch_id, data in enumerate(val_reader()):
        video_feat = np.array([item[0] for item in data]).astype(DATATYPE)
        gt_iou_map = np.array([item[1] for item in data]).astype(DATATYPE)
        gt_start = np.array([item[2] for item in data]).astype(DATATYPE)
        gt_end = np.array([item[3] for item in data]).astype(DATATYPE)

        x_data = to_variable(video_feat)
        gt_iou_map = to_variable(gt_iou_map)
        gt_start = to_variable(gt_start)
        gt_end = to_variable(gt_end)
        gt_iou_map.stop_gradient = True
        gt_start.stop_gradient = True
        gt_end.stop_gradient = True

        pred_bm, pred_start, pred_end = model(x_data)

        loss, tem_loss, pem_reg_loss, pem_cls_loss = bmn_loss_func(
625 626
            pred_bm, pred_start, pred_end, gt_iou_map, gt_start, gt_end, args
        )
627
        avg_loss = paddle.mean(loss)
628 629

        loss_data += [
630 631 632
            avg_loss.numpy()[0],
            tem_loss.numpy()[0],
            pem_reg_loss.numpy()[0],
633
            pem_cls_loss.numpy()[0],
634 635
        ]

636 637 638 639 640 641 642 643 644
        print(
            '[VALID] iter {} '.format(batch_id)
            + '\tLoss = {}, \ttem_loss = {}, \tpem_reg_loss = {}, \tpem_cls_loss = {}'.format(
                '%f' % avg_loss.numpy()[0],
                '%f' % tem_loss.numpy()[0],
                '%f' % pem_reg_loss.numpy()[0],
                '%f' % pem_cls_loss.numpy()[0],
            )
        )
645 646 647 648 649 650 651 652 653

        if batch_id == args.valid_batch_num:
            break
    return loss_data


class TestTrain(unittest.TestCase):
    def setUp(self):
        self.args = Args()
654 655 656
        self.place = (
            fluid.CPUPlace()
            if not fluid.is_compiled_with_cuda()
657
            else fluid.CUDAPlace(0)
658
        )
659

660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686
        self.temp_dir = tempfile.TemporaryDirectory()
        self.model_save_dir = os.path.join(self.temp_dir.name, 'inference')
        self.model_save_prefix = os.path.join(self.model_save_dir, 'bmn')
        self.model_filename = "bmn" + INFER_MODEL_SUFFIX
        self.params_filename = "bmn" + INFER_PARAMS_SUFFIX
        self.dy_param_path = os.path.join(self.temp_dir.name, 'bmn_dy_param')

    def tearDown(self):
        self.temp_dir.cleanup()

    def train_bmn(self, args, place, to_static):
        program_translator.enable(to_static)
        loss_data = []

        with fluid.dygraph.guard(place):
            paddle.seed(SEED)
            paddle.framework.random._manual_program_seed(SEED)
            global local_random
            local_random = np.random.RandomState(SEED)

            bmn = BMN(args)
            adam = optimizer(args, parameter_list=bmn.parameters())

            train_reader = fake_data_reader(args, 'train')

            for epoch in range(args.epoch):
                for batch_id, data in enumerate(train_reader()):
687 688 689 690 691 692 693 694 695 696 697 698
                    video_feat = np.array([item[0] for item in data]).astype(
                        DATATYPE
                    )
                    gt_iou_map = np.array([item[1] for item in data]).astype(
                        DATATYPE
                    )
                    gt_start = np.array([item[2] for item in data]).astype(
                        DATATYPE
                    )
                    gt_end = np.array([item[3] for item in data]).astype(
                        DATATYPE
                    )
699 700 701 702 703 704 705 706 707 708 709 710

                    x_data = to_variable(video_feat)
                    gt_iou_map = to_variable(gt_iou_map)
                    gt_start = to_variable(gt_start)
                    gt_end = to_variable(gt_end)
                    gt_iou_map.stop_gradient = True
                    gt_start.stop_gradient = True
                    gt_end.stop_gradient = True

                    pred_bm, pred_start, pred_end = bmn(x_data)

                    loss, tem_loss, pem_reg_loss, pem_cls_loss = bmn_loss_func(
711 712 713 714 715 716 717 718
                        pred_bm,
                        pred_start,
                        pred_end,
                        gt_iou_map,
                        gt_start,
                        gt_end,
                        args,
                    )
719
                    avg_loss = paddle.mean(loss)
720 721 722 723 724 725

                    avg_loss.backward()
                    adam.minimize(avg_loss)
                    bmn.clear_gradients()
                    # log loss data to verify correctness
                    loss_data += [
726 727 728
                        avg_loss.numpy()[0],
                        tem_loss.numpy()[0],
                        pem_reg_loss.numpy()[0],
729
                        pem_cls_loss.numpy()[0],
730 731
                    ]

732 733 734 735 736 737 738 739 740 741 742 743
                    if args.log_interval > 0 and (
                        batch_id % args.log_interval == 0
                    ):
                        print(
                            '[TRAIN] Epoch {}, iter {} '.format(epoch, batch_id)
                            + '\tLoss = {}, \ttem_loss = {}, \tpem_reg_loss = {}, \tpem_cls_loss = {}'.format(
                                '%f' % avg_loss.numpy()[0],
                                '%f' % tem_loss.numpy()[0],
                                '%f' % pem_reg_loss.numpy()[0],
                                '%f' % pem_cls_loss.numpy()[0],
                            )
                        )
744 745 746 747 748 749 750 751 752 753

                    # validation
                    if batch_id % args.valid_interval == 0 and batch_id > 0:
                        bmn.eval()
                        val_loss_data = val_bmn(bmn, args)
                        bmn.train()
                        loss_data += val_loss_data

                    if batch_id == args.train_batch_num:
                        if to_static:
754
                            paddle.jit.save(bmn, self.model_save_prefix)
755
                        else:
756 757 758
                            fluid.dygraph.save_dygraph(
                                bmn.state_dict(), self.dy_param_path
                            )
759 760 761
                        break
            return np.array(loss_data)

762 763
    def test_train(self):

764 765
        static_res = self.train_bmn(self.args, self.place, to_static=True)
        dygraph_res = self.train_bmn(self.args, self.place, to_static=False)
766 767 768 769 770
        np.testing.assert_allclose(
            dygraph_res,
            static_res,
            rtol=1e-05,
            err_msg='dygraph_res: {},\n static_res: {}'.format(
771
                dygraph_res[~np.isclose(dygraph_res, static_res)],
772 773 774 775
                static_res[~np.isclose(dygraph_res, static_res)],
            ),
            atol=1e-8,
        )
776 777 778 779 780 781 782 783 784 785 786 787

        # Prediction needs trained models, so put `test_predict` at last of `test_train`
        self.verify_predict()

    def verify_predict(self):
        args = Args()
        args.batch_size = 1  # change batch_size
        test_reader = fake_data_reader(args, 'test')
        for batch_id, data in enumerate(test_reader()):
            video_data = np.array([item[0] for item in data]).astype(DATATYPE)
            static_pred_res = self.predict_static(video_data)
            dygraph_pred_res = self.predict_dygraph(video_data)
788
            dygraph_jit_pred_res = self.predict_dygraph_jit(video_data)
789
            predictor_pred_res = self.predict_analysis_inference(video_data)
790

791
            for dy_res, st_res, dy_jit_res, predictor_res in zip(
792 793 794 795 796
                dygraph_pred_res,
                static_pred_res,
                dygraph_jit_pred_res,
                predictor_pred_res,
            ):
797 798 799 800 801
                np.testing.assert_allclose(
                    st_res,
                    dy_res,
                    rtol=1e-05,
                    err_msg='dygraph_res: {},\n static_res: {}'.format(
802
                        dy_res[~np.isclose(st_res, dy_res)],
803 804 805 806
                        st_res[~np.isclose(st_res, dy_res)],
                    ),
                    atol=1e-8,
                )
807 808 809 810 811
                np.testing.assert_allclose(
                    st_res,
                    dy_jit_res,
                    rtol=1e-05,
                    err_msg='dygraph_jit_res: {},\n static_res: {}'.format(
812
                        dy_jit_res[~np.isclose(st_res, dy_jit_res)],
813 814 815 816
                        st_res[~np.isclose(st_res, dy_jit_res)],
                    ),
                    atol=1e-8,
                )
817 818 819 820 821
                np.testing.assert_allclose(
                    st_res,
                    predictor_res,
                    rtol=1e-05,
                    err_msg='dygraph_jit_res: {},\n static_res: {}'.format(
822
                        predictor_res[~np.isclose(st_res, predictor_res)],
823 824 825 826
                        st_res[~np.isclose(st_res, predictor_res)],
                    ),
                    atol=1e-8,
                )
827 828 829 830 831 832 833
            break

    def predict_dygraph(self, data):
        program_translator.enable(False)
        with fluid.dygraph.guard(self.place):
            bmn = BMN(self.args)
            # load dygraph trained parameters
834
            model_dict, _ = fluid.load_dygraph(self.dy_param_path + ".pdparams")
835 836 837 838 839 840 841 842 843 844
            bmn.set_dict(model_dict)
            bmn.eval()

            x = to_variable(data)
            pred_res = bmn(x)
            pred_res = [var.numpy() for var in pred_res]

            return pred_res

    def predict_static(self, data):
845
        paddle.enable_static()
846 847
        exe = fluid.Executor(self.place)
        # load inference model
848 849 850 851 852 853 854 855 856 857 858 859 860 861 862
        [
            inference_program,
            feed_target_names,
            fetch_targets,
        ] = fluid.io.load_inference_model(
            self.model_save_dir,
            executor=exe,
            model_filename=self.model_filename,
            params_filename=self.params_filename,
        )
        pred_res = exe.run(
            inference_program,
            feed={feed_target_names[0]: data},
            fetch_list=fetch_targets,
        )
863 864 865

        return pred_res

866 867
    def predict_dygraph_jit(self, data):
        with fluid.dygraph.guard(self.place):
868
            bmn = paddle.jit.load(self.model_save_prefix)
869 870 871 872 873 874 875 876
            bmn.eval()

            x = to_variable(data)
            pred_res = bmn(x)
            pred_res = [var.numpy() for var in pred_res]

            return pred_res

877
    def predict_analysis_inference(self, data):
878 879 880 881 882 883
        output = PredictorTools(
            self.model_save_dir,
            self.model_filename,
            self.params_filename,
            [data],
        )
884 885 886
        out = output()
        return out

887 888

if __name__ == "__main__":
889 890
    with fluid.framework._test_eager_guard():
        unittest.main()