Python mxnet.nd.expand_dims() Examples

def extract_pairwise_multi_position_embedding_nd(position_mat, feat_dim, wave_length=1000):
    """ Extract multi-class position embedding

    Args:
        position_mat: [num_fg_classes, num_rois, num_rois, 4]
        feat_dim: dimension of embedding feature
        wave_length:

    Returns:
        embedding: [num_fg_classes, num_rois, num_rois, feat_dim]
    """
    feat_range = nd.arange(0, feat_dim / 8)
    dim_mat = nd.broadcast_power(lhs=nd.full((1,), wave_length),
                                     rhs=(8. / feat_dim) * feat_range)
    dim_mat = nd.Reshape(dim_mat, shape=(1, 1, 1, 1, -1))
    position_mat = nd.expand_dims(100.0 * position_mat, axis=4)
    div_mat = nd.broadcast_div(lhs=position_mat, rhs=dim_mat)
    sin_mat = nd.sin(data=div_mat)
    cos_mat = nd.cos(data=div_mat)
    # embedding, [num_fg_classes, num_rois, num_rois, 4, feat_dim/4]
    embedding = nd.concat(sin_mat, cos_mat, dim=4)
    embedding = nd.Reshape(embedding, shape=(0, 0, 0, feat_dim))
    return embedding 

def extract_pairwise_multi_position_embedding_nd(position_mat, feat_dim, wave_length=1000):
    """ Extract multi-class position embedding

    Args:
        position_mat: [num_fg_classes, num_rois, num_rois, 4]
        feat_dim: dimension of embedding feature
        wave_length:

    Returns:
        embedding: [num_fg_classes, num_rois, num_rois, feat_dim]
    """
    feat_range = nd.arange(0, feat_dim / 8)
    dim_mat = nd.broadcast_power(lhs=nd.full((1,), wave_length),
                                     rhs=(8. / feat_dim) * feat_range)
    dim_mat = nd.Reshape(dim_mat, shape=(1, 1, 1, 1, -1))
    position_mat = nd.expand_dims(100.0 * position_mat, axis=4)
    div_mat = nd.broadcast_div(lhs=position_mat, rhs=dim_mat)
    sin_mat = nd.sin(data=div_mat)
    cos_mat = nd.cos(data=div_mat)
    # embedding, [num_fg_classes, num_rois, num_rois, 4, feat_dim/4]
    embedding = nd.concat(sin_mat, cos_mat, dim=4)
    embedding = nd.Reshape(embedding, shape=(0, 0, 0, feat_dim))
    return embedding 

def __init__(self, dist, edge, hidden_size, prefix=None):
        super(Graph, self).__init__(prefix=prefix)
        self.dist = dist
        self.edge = edge
        self.hidden_size = hidden_size

        # create graph
        self.num_nodes = n = self.dist.shape[0]
        src, dst, dist = [], [], []
        for i in range(n):
            for j in edge[i]:
                src.append(j)
                dst.append(i)
                dist.append(self.dist[j, i])

        self.src = src
        self.dst = dst
        self.dist = mx.nd.expand_dims(mx.nd.array(dist), axis=1)
        self.ctx = []
        self.graph_on_ctx = []

        self.init_model() 

def Route(self, x):
        # b_mat = nd.repeat(self.b_mat.data(), repeats=x.shape[0], axis=0)#nd.stop_gradient(nd.repeat(self.b_mat.data(), repeats=x.shape[0], axis=0))
        b_mat = nd.zeros((x.shape[0],1,self.num_cap, self.num_locations), ctx=x.context)
        x_expand = nd.expand_dims(nd.expand_dims(x, axis=2),2)
        w_expand = nd.repeat(nd.expand_dims(self.w_ij.data(x.context),axis=0), repeats=x.shape[0], axis=0)
        u_ = w_expand*x_expand
        # u_ = nd.abs(w_expand - x_expand)
        u = nd.sum(u_, axis = 1)
        u_no_gradient = nd.stop_gradient(u)
        for i in range(self.route_num):
            c_mat = nd.softmax(b_mat, axis=2)
            if i == self.route_num -1:
                s = nd.sum(u * c_mat, axis=-1)
            else:
                s = nd.sum(u_no_gradient * c_mat, axis=-1)
            v = squash(s, 1)
            v1 = nd.expand_dims(v, axis=-1)
            if i != self.route_num - 1:
                update_term = nd.sum(u_no_gradient*v1, axis=1, keepdims=True)
                b_mat = b_mat + update_term
        return v 

def Route(self, x):
        # print x.context
        b_mat = nd.zeros((x.shape[0],1,self.num_cap, self.num_locations), ctx=x.context)
        x_expand = nd.expand_dims(nd.expand_dims(x, axis=2),2)
        w_expand = nd.repeat(nd.expand_dims(self.w_ij.data(x.context),axis=0), repeats=x.shape[0], axis=0)
        u_ = w_expand*x_expand
        u = nd.sum(u_, axis = 1)
        # u_ = nd.square(w_expand - x_expand)
        # u = -nd.sum(u_, axis = 1)
        u_no_gradient = nd.stop_gradient(u)
        for i in range(self.route_num):
            # c_mat = nd.softmax(b_mat, axis=2)
            c_mat = nd.sigmoid(b_mat)
            if i == self.route_num -1:
                s = nd.sum(u * c_mat, axis=-1)
            else:
                s = nd.sum(u_no_gradient * c_mat, axis=-1)
            v = squash(s, 1)
            if i != self.route_num - 1:
                v1 = nd.expand_dims(v, axis=-1)
                update_term = nd.sum(u_no_gradient*v1, axis=1, keepdims=True)
                b_mat = b_mat + update_term
                # b_mat = update_term
            # else:
            #    v = s
        return v 

def extract_multi_position_matrix_nd(bbox):
    bbox = nd.transpose(bbox, axes=(1, 0, 2))
    xmin, ymin, xmax, ymax = nd.split(data=bbox, num_outputs=4, axis=2)
    # [num_fg_classes, num_boxes, 1]
    bbox_width = xmax - xmin + 1.
    bbox_height = ymax - ymin + 1.
    center_x = 0.5 * (xmin + xmax)
    center_y = 0.5 * (ymin + ymax)
    # [num_fg_classes, num_boxes, num_boxes]
    delta_x = nd.broadcast_minus(lhs=center_x, 
        rhs=nd.transpose(center_x, axes=(0, 2, 1)))
    delta_x = nd.broadcast_div(delta_x, bbox_width)
    delta_x = nd.log(nd.maximum(nd.abs(delta_x), 1e-3))

    delta_y = nd.broadcast_minus(lhs=center_y,
        rhs=nd.transpose(center_y, axes=(0, 2, 1)))
    delta_y = nd.broadcast_div(delta_y, bbox_height)
    delta_y = nd.log(nd.maximum(nd.abs(delta_y), 1e-3))

    delta_width = nd.broadcast_div(lhs=bbox_width, 
        rhs=nd.transpose(bbox_width, axes=(0, 2, 1)))
    delta_width = nd.log(delta_width)

    delta_height = nd.broadcast_div(lhs=bbox_height,
        rhs=nd.transpose(bbox_height, axes=(0, 2, 1)))
    delta_height = nd.log(delta_height)
    concat_list = [delta_x, delta_y, delta_width, delta_height]
    for idx, sym in enumerate(concat_list):
        concat_list[idx] = nd.expand_dims(sym, axis=3)
    position_matrix = nd.concat(*concat_list, dim=3)
    return position_matrix 

def extract_rank_embedding_nd(rank_dim, feat_dim, wave_length=1000):
    rank_range = nd.arange(0, rank_dim)
    feat_range = nd.arange(0, feat_dim / 2)
    dim_mat = nd.broadcast_power(lhs=nd.full((1,), wave_length),
                                     rhs=(2. / feat_dim) * feat_range)
    dim_mat = nd.Reshape(dim_mat, shape=(1, -1))
    rank_mat = nd.expand_dims(rank_range, axis=1)
    div_mat = nd.broadcast_div(lhs=rank_mat, rhs=dim_mat)
    sin_mat = nd.sin(data=div_mat)
    cos_mat = nd.cos(data=div_mat)
    embedding = nd.concat(sin_mat, cos_mat, dim=1)
    return embedding 

def forward_single(self, feature, data, begin_state):
        # add a temporal axis
        data = nd.expand_dims(data, axis=2)

        # unroll
        data, state = self(feature, data, begin_state)

        # remove the temporal axis
        data = nd.mean(data, axis=2)

        return data, state 

def forward_single(self, feature, data, begin_state):
        # add a temporal axis
        data = nd.expand_dims(data, axis=2)

        # unroll
        data, state = self(feature, data, begin_state)

        # remove the temporal axis
        data = nd.mean(data, axis=2)

        return data, state 

def msg_edge(self, edge):
        state = nd.concat(edge.src['state'], edge.dst['state'], dim=-1)
        ctx = state.context

        alpha = nd.LeakyReLU(nd.dot(state, self.weight.data(ctx)))

        dist = edge.data['dist']
        while len(dist.shape) < len(alpha.shape):
            dist = nd.expand_dims(dist, axis=-1)

        alpha = alpha * dist 
        return { 'alpha': alpha, 'state': edge.src['state'] } 

def forward(self, x):
        conv_out = nd.expand_dims(self.cap(x), axis=2)
        conv_out = conv_out.reshape((0,-1,self.num_cap,0,0))
        conv_out  = squash(conv_out, 1)
        return conv_out 

def route(self, u):
        b_mat = nd.zeros((u.shape[0], self.num_cap_in, self.num_cap, 1, u.shape[4], u.shape[5]), ctx=u.context)
        for i in range(self.route_num):
            c_mat = nd.softmax(b_mat, axis=2)
            s = nd.sum(u * c_mat, axis=1)
            v = squash(s, 2)
            if i != self.route_num - 1:
                v1 = nd.expand_dims(v, axis=1)
                update_term = nd.sum(u*v1, axis=3, keepdims=True)
                b_mat = b_mat + update_term
        return v 

def forward(self, x):
        conv_out = nd.expand_dims(self.cap(x), axis=2)
        # conv_out = nd.expand_dims(self.bn(self.cap(x)), axis=2)
        conv_out = conv_out.reshape((0,self.num_cap,-1,0,0))
        conv_out = squash(conv_out, 2)
        # print conv_out.shape
        return conv_out 

def forward(self, x):
        x1 = nd.expand_dims(x, axis=self.axes)
        return x1 

def extract_multi_position_matrix_nd(bbox):
    bbox = nd.transpose(bbox, axes=(1, 0, 2))
    xmin, ymin, xmax, ymax = nd.split(data=bbox, num_outputs=4, axis=2)
    # [num_fg_classes, num_boxes, 1]
    bbox_width = xmax - xmin + 1.
    bbox_height = ymax - ymin + 1.
    center_x = 0.5 * (xmin + xmax)
    center_y = 0.5 * (ymin + ymax)
    # [num_fg_classes, num_boxes, num_boxes]
    delta_x = nd.broadcast_minus(lhs=center_x, 
        rhs=nd.transpose(center_x, axes=(0, 2, 1)))
    delta_x = nd.broadcast_div(delta_x, bbox_width)
    delta_x = nd.log(nd.maximum(nd.abs(delta_x), 1e-3))

    delta_y = nd.broadcast_minus(lhs=center_y,
        rhs=nd.transpose(center_y, axes=(0, 2, 1)))
    delta_y = nd.broadcast_div(delta_y, bbox_height)
    delta_y = nd.log(nd.maximum(nd.abs(delta_y), 1e-3))

    delta_width = nd.broadcast_div(lhs=bbox_width, 
        rhs=nd.transpose(bbox_width, axes=(0, 2, 1)))
    delta_width = nd.log(delta_width)

    delta_height = nd.broadcast_div(lhs=bbox_height,
        rhs=nd.transpose(bbox_height, axes=(0, 2, 1)))
    delta_height = nd.log(delta_height)
    concat_list = [delta_x, delta_y, delta_width, delta_height]
    for idx, sym in enumerate(concat_list):
        concat_list[idx] = nd.expand_dims(sym, axis=3)
    position_matrix = nd.concat(*concat_list, dim=3)
    return position_matrix 

def extract_rank_embedding_nd(rank_dim, feat_dim, wave_length=1000):
    rank_range = nd.arange(0, rank_dim)
    feat_range = nd.arange(0, feat_dim / 2)
    dim_mat = nd.broadcast_power(lhs=nd.full((1,), wave_length),
                                     rhs=(2. / feat_dim) * feat_range)
    dim_mat = nd.Reshape(dim_mat, shape=(1, -1))
    rank_mat = nd.expand_dims(rank_range, axis=1)
    div_mat = nd.broadcast_div(lhs=rank_mat, rhs=dim_mat)
    sin_mat = nd.sin(data=div_mat)
    cos_mat = nd.cos(data=div_mat)
    embedding = nd.concat(sin_mat, cos_mat, dim=1)
    return embedding 

def init_weights(self, ctx):
        """Initial I3D network with its 2D pretrained weights."""

        self.first_stage.initialize(ctx=ctx)
        self.res_layers.initialize(ctx=ctx)
        self.head.initialize(ctx=ctx)

        if self.pretrained_base and not self.pretrained:
            if self.depth == 50:
                resnet2d = resnet50_v1b(pretrained=True)
            elif self.depth == 101:
                resnet2d = resnet101_v1b(pretrained=True)
            else:
                print('No such 2D pre-trained network of depth %d.' % (self.depth))

            weights2d = resnet2d.collect_params()
            if self.nonlocal_cfg is None:
                weights3d = self.collect_params()
            else:
                train_params_list = []
                raw_params = self.collect_params()
                for raw_name in raw_params.keys():
                    if 'nonlocal' in raw_name:
                        continue
                    train_params_list.append(raw_name)
                init_patterns = '|'.join(train_params_list)
                weights3d = self.collect_params(init_patterns)
            assert len(weights2d.keys()) == len(weights3d.keys()), 'Number of parameters should be same.'

            dict2d = {}
            for key_id, key_name in enumerate(weights2d.keys()):
                dict2d[key_id] = key_name

            dict3d = {}
            for key_id, key_name in enumerate(weights3d.keys()):
                dict3d[key_id] = key_name

            dict_transform = {}
            for key_id, key_name in dict3d.items():
                dict_transform[dict2d[key_id]] = key_name

            cnt = 0
            for key2d, key3d in dict_transform.items():
                if 'conv' in key3d:
                    temporal_dim = weights3d[key3d].shape[2]
                    temporal_2d = nd.expand_dims(weights2d[key2d].data(), axis=2)
                    inflated_2d = nd.broadcast_to(temporal_2d, shape=[0, 0, temporal_dim, 0, 0]) / temporal_dim
                    assert inflated_2d.shape == weights3d[key3d].shape, 'the shape of %s and %s does not match. ' % (key2d, key3d)
                    weights3d[key3d].set_data(inflated_2d)
                    cnt += 1
                    print('%s is done with shape: ' % (key3d), weights3d[key3d].shape)
                if 'batchnorm' in key3d:
                    assert weights2d[key2d].shape == weights3d[key3d].shape, 'the shape of %s and %s does not match. ' % (key2d, key3d)
                    weights3d[key3d].set_data(weights2d[key2d].data())
                    cnt += 1
                    print('%s is done with shape: ' % (key3d), weights3d[key3d].shape)
                if 'dense' in key3d:
                    cnt += 1
                    print('%s is skipped with shape: ' % (key3d), weights3d[key3d].shape)

            assert cnt == len(weights2d.keys()), 'Not all parameters have been ported, check the initialization.' 

def forward(self, feature, label, begin_states, is_training):
        ''' Decode the hidden states to a temporal sequence.

        Parameters
        ----------
        feature: a NDArray with shape [n, d].
        label: a NDArray with shape [n, b, t, d].
        begin_states: a list of hidden states (list of hidden units with shape [n, b, d]) of RNNs.
        is_training: bool
        
        Returns
        -------
            outputs: the prediction, which is a NDArray with shape [n, b, t, d]
        '''
        ctx = label.context

        num_nodes, batch_size, seq_len, _ = label.shape 
        aux = label[:,:,:, self.output_dim:] # [n,b,t,d]
        label = label[:,:,:, :self.output_dim] # [n,b,t,d]
        
        go = nd.zeros(shape=(num_nodes, batch_size, self.input_dim), ctx=ctx)
        output, states = [], begin_states

        for i in range(seq_len):
            # get next input
            if i == 0: data = go
            else:
                prev = nd.concat(output[i - 1], aux[:,:,i - 1], dim=-1)
                truth = nd.concat(label[:,:,i - 1], aux[:,:,i - 1], dim=-1)
                if is_training and self.use_sampling: value = self.sampling()
                else: value = 0
                data = value * truth + (1 - value) * prev

            # unroll 1 step
            for depth, cell in enumerate(self.cells):
                data, states[depth] = cell.forward_single(feature, data, states[depth])
                if self.graphs[depth] is not None:
                    _data = 0
                    for g in self.graphs[depth]:
                        _data = _data + g(data, feature)
                    data = _data / len(self.graphs[depth])

            # append feature to output
            _feature = nd.expand_dims(feature, axis=1) # [n, 1, d]
            _feature = nd.broadcast_to(_feature, shape=(0, batch_size, 0)) # [n, b, d]
            data = nd.concat(data, _feature, dim=-1) # [n, b, t, d]

            # proj output to prediction
            data = nd.reshape(data, shape=(num_nodes * batch_size, -1))
            data = self.proj(data)
            data = nd.reshape(data, shape=(num_nodes, batch_size, -1))
            
            output.append(data)

        output = nd.stack(*output, axis=2)
        return output 

def forward(self, feature, label, begin_states, is_training):
        ''' Decode the hidden states to a temporal sequence.

        Parameters
        ----------
        feature: a NDArray with shape [n, d].
        label: a NDArray with shape [n, b, t, d].
        begin_states: a list of hidden states (list of hidden units with shape [n, b, d]) of RNNs.
        is_training: bool
        
        Returns
        -------
            outputs: the prediction, which is a NDArray with shape [n, b, t, d]
        '''
        ctx = label.context

        num_nodes, batch_size, seq_len, _ = label.shape 
        aux = label[:,:,:, self.output_dim:] # [n,b,t,d]
        label = label[:,:,:, :self.output_dim] # [n,b,t,d]
        
        go = nd.zeros(shape=(num_nodes, batch_size, self.input_dim), ctx=ctx)
        output, states = [], begin_states

        for i in range(seq_len):
            # get next input
            if i == 0: data = go
            else:
                prev = nd.concat(output[i - 1], aux[:,:,i - 1], dim=-1)
                truth = nd.concat(label[:,:,i - 1], aux[:,:,i - 1], dim=-1)
                if is_training and self.use_sampling: value = self.sampling()
                else: value = 0
                data = value * truth + (1 - value) * prev

            # unroll 1 step
            for depth, cell in enumerate(self.cells):
                data, states[depth] = cell.forward_single(feature, data, states[depth])
                if self.graphs[depth] is not None:
                    _data = data
                    for g in self.graphs[depth]:
                        _data = _data + g(data, feature)
                    data = _data

            # append feature to output
            _feature = nd.expand_dims(feature, axis=1) # [n, 1, d]
            _feature = nd.broadcast_to(_feature, shape=(0, batch_size, 0)) # [n, b, d]
            data = nd.concat(data, _feature, dim=-1) # [n, b, t, d]

            # proj output to prediction
            data = nd.reshape(data, shape=(num_nodes * batch_size, -1))
            data = self.proj(data)
            data = nd.reshape(data, shape=(num_nodes, batch_size, -1))
            
            output.append(data)

        output = nd.stack(*output, axis=2)
        return output 

def refine_bbox_nd(bbox, bbox_delta, im_info=None, means=None, stds=None):

    xmin, ymin, xmax, ymax = nd.split(data=bbox, num_outputs=4, axis=1)
    bbox_width = xmax - xmin + 1.
    bbox_height = ymax - ymin + 1.
    center_x = 0.5 * (xmin + xmax)
    center_y = 0.5 * (ymin + ymax)

    bbox_delta_reshape = nd.Reshape(data=bbox_delta, shape=(0, -1, 4))
    dx, dy, dw, dh = nd.split(data=bbox_delta_reshape, 
        num_outputs=4, axis=2, squeeze_axis=1)
    if (means is not None) and (stds is not None):
        dx = dx * stds[0] + means[0]
        dy = dy * stds[1] + means[1]
        dw = dw * stds[2] + means[2]
        dh = dh * stds[3] + means[3]

    refine_center_x = nd.broadcast_add(lhs=center_x,
        rhs=nd.broadcast_mul(lhs=bbox_width, rhs=dx))
    refine_center_y = nd.broadcast_add(lhs=center_y,
        rhs=nd.broadcast_mul(lhs=bbox_height, rhs=dy))
    refined_width = nd.broadcast_mul(lhs=bbox_width, rhs=nd.exp(dw))
    refined_height = nd.broadcast_mul(lhs=bbox_height, rhs=nd.exp(dh))
    w_offset = 0.5 * (refined_width - 1.)
    h_offset = 0.5 * (refined_height - 1.)
    refined_xmin = nd.expand_dims(refine_center_x - w_offset, axis=1)
    refined_ymin = nd.expand_dims(refine_center_y - h_offset, axis=1)
    refined_xmax = nd.expand_dims(refine_center_x + w_offset, axis=1)
    refined_ymax = nd.expand_dims(refine_center_y + h_offset, axis=1)

    refined_bbox = nd.concat(refined_xmin, refined_ymin, refined_xmax, refined_ymax, dim=1)
    if im_info is not None:
        # assume im_info [[height, width, scale]] with shape (1,3)
        im_hw = nd.slice_axis(im_info, axis=1, begin=0, end=2)
        im_wh = nd.reverse(im_hw, axis=1)
        im_wh = im_wh - 1.
        im_wh = nd.tile(data=im_wh, reps=(1, 2))
        im_wh = nd.Reshape(im_wh, shape=(1, 4, 1))
        refined_bbox = nd.broadcast_minimum(lhs=refined_bbox, rhs=im_wh)
        refined_bbox = nd.broadcast_maximum(lhs=refined_bbox,
            rhs=nd.zeros_like(refined_bbox))
    # print refined_bbox.debug_str()
    return refined_bbox 

def refine_bbox_nd(bbox, bbox_delta, im_info=None, means=None, stds=None):

    xmin, ymin, xmax, ymax = nd.split(data=bbox, num_outputs=4, axis=1)
    bbox_width = xmax - xmin + 1.
    bbox_height = ymax - ymin + 1.
    center_x = 0.5 * (xmin + xmax)
    center_y = 0.5 * (ymin + ymax)

    bbox_delta_reshape = nd.Reshape(data=bbox_delta, shape=(0, -1, 4))
    dx, dy, dw, dh = nd.split(data=bbox_delta_reshape, 
        num_outputs=4, axis=2, squeeze_axis=1)
    if (means is not None) and (stds is not None):
        dx = dx * stds[0] + means[0]
        dy = dy * stds[1] + means[1]
        dw = dw * stds[2] + means[2]
        dh = dh * stds[3] + means[3]

    refine_center_x = nd.broadcast_add(lhs=center_x,
        rhs=nd.broadcast_mul(lhs=bbox_width, rhs=dx))
    refine_center_y = nd.broadcast_add(lhs=center_y,
        rhs=nd.broadcast_mul(lhs=bbox_height, rhs=dy))
    refined_width = nd.broadcast_mul(lhs=bbox_width, rhs=nd.exp(dw))
    refined_height = nd.broadcast_mul(lhs=bbox_height, rhs=nd.exp(dh))
    w_offset = 0.5 * (refined_width - 1.)
    h_offset = 0.5 * (refined_height - 1.)
    refined_xmin = nd.expand_dims(refine_center_x - w_offset, axis=1)
    refined_ymin = nd.expand_dims(refine_center_y - h_offset, axis=1)
    refined_xmax = nd.expand_dims(refine_center_x + w_offset, axis=1)
    refined_ymax = nd.expand_dims(refine_center_y + h_offset, axis=1)

    refined_bbox = nd.concat(refined_xmin, refined_ymin, refined_xmax, refined_ymax, dim=1)
    if im_info is not None:
        # assume im_info [[height, width, scale]] with shape (1,3)
        im_hw = nd.slice_axis(im_info, axis=1, begin=0, end=2)
        im_wh = nd.reverse(im_hw, axis=1)
        im_wh = im_wh - 1.
        im_wh = nd.tile(data=im_wh, reps=(1, 2))
        im_wh = nd.Reshape(im_wh, shape=(1, 4, 1))
        refined_bbox = nd.broadcast_minimum(lhs=refined_bbox, rhs=im_wh)
        refined_bbox = nd.broadcast_maximum(lhs=refined_bbox,
            rhs=nd.zeros_like(refined_bbox))
    # print refined_bbox.debug_str()
    return refined_bbox