TypeScript tfjs-core.concat函数代码示例

export function residualDown(x: tf.Tensor4D, params: ResidualLayerParams): tf.Tensor4D {
  let out = convDown(x, params.conv1)
  out = convNoRelu(out, params.conv2)

  let pooled = tf.avgPool(x, 2, 2, 'valid') as tf.Tensor4D
  const zeros = tf.zeros<tf.Rank.R4>(pooled.shape)
  const isPad = pooled.shape[3] !== out.shape[3]
  const isAdjustShape = pooled.shape[1] !== out.shape[1] || pooled.shape[2] !== out.shape[2]

  if (isAdjustShape) {
    const padShapeX = [...out.shape] as [number, number, number, number]
    padShapeX[1] = 1
    const zerosW = tf.zeros<tf.Rank.R4>(padShapeX)
    out = tf.concat([out, zerosW], 1)

    const padShapeY = [...out.shape] as [number, number, number, number]
    padShapeY[2] = 1
    const zerosH = tf.zeros<tf.Rank.R4>(padShapeY)
    out = tf.concat([out, zerosH], 2)
  }

  pooled = isPad ? tf.concat([pooled, zeros], 3) : pooled
  out = tf.add(pooled, out) as tf.Tensor4D

  out = tf.relu(out)
  return out
}

  return tf.tidy(() => {

    const [height, width] = imgTensor.shape.slice(1)
    if (height === width) {
      return imgTensor
    }

    const dimDiff = Math.abs(height - width)
    const paddingAmount = Math.round(dimDiff * (isCenterImage ? 0.5 : 1))
    const paddingAxis = height > width ? 2 : 1

    const createPaddingTensor = (paddingAmount: number): tf.Tensor => {
      const paddingTensorShape = imgTensor.shape.slice()
      paddingTensorShape[paddingAxis] = paddingAmount
      return tf.fill(paddingTensorShape, 0)
    }

    const paddingTensorAppend = createPaddingTensor(paddingAmount)
    const remainingPaddingAmount = dimDiff - paddingTensorAppend.shape[paddingAxis]

    const paddingTensorPrepend = isCenterImage && remainingPaddingAmount
      ? createPaddingTensor(remainingPaddingAmount)
      : null

    const tensorsToStack = [
      paddingTensorPrepend,
      imgTensor,
      paddingTensorAppend
    ]
      .filter(t => t !== null) as tf.Tensor4D[]
    return tf.concat(tensorsToStack, paddingAxis)
  })

  return tf.tidy(() => {
    const avg_r = tf.fill([1, 150, 150, 1], 122.782);
    const avg_g = tf.fill([1, 150, 150, 1], 117.001);
    const avg_b = tf.fill([1, 150, 150, 1], 104.298);
    const avg_rgb = tf.concat([avg_r, avg_g, avg_b], 3)

    return tf.div(tf.sub(x, avg_rgb), tf.scalar(256))
  })

  return tf.tidy(() => {

    const conv0 = pointwiseConvLayer(x, params.conv_0_params, [1, 1])
    const conv1 = pointwiseConvLayer(conv0, params.conv_1_params, [2, 2])
    const conv2 = pointwiseConvLayer(conv1, params.conv_2_params, [1, 1])
    const conv3 = pointwiseConvLayer(conv2, params.conv_3_params, [2, 2])
    const conv4 = pointwiseConvLayer(conv3, params.conv_4_params, [1, 1])
    const conv5 = pointwiseConvLayer(conv4, params.conv_5_params, [2, 2])
    const conv6 = pointwiseConvLayer(conv5, params.conv_6_params, [1, 1])
    const conv7 = pointwiseConvLayer(conv6, params.conv_7_params, [2, 2])

    const boxPrediction0 = boxPredictionLayer(conv11, params.box_predictor_0_params)
    const boxPrediction1 = boxPredictionLayer(x, params.box_predictor_1_params)
    const boxPrediction2 = boxPredictionLayer(conv1, params.box_predictor_2_params)
    const boxPrediction3 = boxPredictionLayer(conv3, params.box_predictor_3_params)
    const boxPrediction4 = boxPredictionLayer(conv5, params.box_predictor_4_params)
    const boxPrediction5 = boxPredictionLayer(conv7, params.box_predictor_5_params)

    const boxPredictions = tf.concat([
      boxPrediction0.boxPredictionEncoding,
      boxPrediction1.boxPredictionEncoding,
      boxPrediction2.boxPredictionEncoding,
      boxPrediction3.boxPredictionEncoding,
      boxPrediction4.boxPredictionEncoding,
      boxPrediction5.boxPredictionEncoding
    ], 1) as tf.Tensor4D

    const classPredictions = tf.concat([
      boxPrediction0.classPrediction,
      boxPrediction1.classPrediction,
      boxPrediction2.classPrediction,
      boxPrediction3.classPrediction,
      boxPrediction4.classPrediction,
      boxPrediction5.classPrediction
    ], 1) as tf.Tensor4D

    return {
      boxPredictions,
      classPredictions
    }
  })

  return tf.tidy(() => {
    if (input instanceof tf.Tensor) {
      const rank = input.shape.length
      if (rank !== 3 && rank !== 4) {
        throw new Error('input tensor must be of rank 3 or 4')
      }

      return (rank === 3 ? input.expandDims(0) : input).toFloat() as tf.Tensor4D
    }

    if (!(input instanceof NetInput)) {
      throw new Error('getImageTensor - expected input to be a tensor or an instance of NetInput')
    }

    return tf.concat(
      input.canvases.map(canvas =>
        tf.fromPixels(canvas).expandDims(0).toFloat()
      )
    ) as tf.Tensor4D
  })

export let executeOp: OpExecutor = (node: Node, tensorMap: NamedTensorsMap,
                                    context: ExecutionContext):
                                       tfc.Tensor[] => {
  switch (node.op) {
    case 'concat': {
      const axis = getParamValue('axis', node, tensorMap, context) as number;
      const inputs =
          getParamValue('tensors', node, tensorMap, context) as tfc.Tensor[];
      return [tfc.concat(inputs, axis)];
    }
    case 'gather': {
      const axis = getParamValue('axis', node, tensorMap, context) as number;
      const input = getParamValue('x', node, tensorMap, context) as tfc.Tensor;
      const indices =
          getParamValue('indices', node, tensorMap, context) as tfc.Tensor1D;
      return [tfc.gather(input, indices, axis)];
    }
    case 'reverse': {
      const axis = getParamValue('axis', node, tensorMap, context) as number;
      const input = getParamValue('x', node, tensorMap, context) as tfc.Tensor;
      return [tfc.reverse(input, axis)];
    }
    case 'slice': {
      // tslint:disable-next-line:no-any
      const begin = getParamValue('begin', node, tensorMap, context) as any;
      // tslint:disable-next-line:no-any
      const size = getParamValue('size', node, tensorMap, context) as any;
      return [tfc.slice(
          getParamValue('x', node, tensorMap, context) as tfc.Tensor, begin,
          size)];
    }
    case 'stridedSlice': {
      const begin =
          getParamValue('begin', node, tensorMap, context) as number[];
      const end = getParamValue('end', node, tensorMap, context) as number[];
      const strides =
          getParamValue('strides', node, tensorMap, context) as number[];
      const beginMask =
          getParamValue('beginMask', node, tensorMap, context) as number;
      const endMask =
          getParamValue('endMask', node, tensorMap, context) as number;
      const tensor = getParamValue('x', node, tensorMap, context) as tfc.Tensor;
      if (begin.length === 1 && tensor.shape.length > 1) {
        for (let i = 1; i < tensor.shape.length; i++) {
          begin.push(0);
          end.push(tensor.shape[i]);
          strides.push(strides[0]);
        }
      }
      return [tfc.stridedSlice(
          tensor, begin, end, strides, beginMask, endMask)];
    }
    case 'stack': {
      return tfc.tidy(() => {
        const axis = getParamValue('axis', node, tensorMap, context) as number;
        const tensors =
            getParamValue('tensors', node, tensorMap, context) as tfc.Tensor[];
        // Reshape the tensors to the first tensor's shape if they don't match.
        const shape = tensors[0].shape;
        const squeezedShape = tensors[0].squeeze().shape;
        const mapped = tensors.map(tensor => {
          const sameShape = tfc.util.arraysEqual(tensor.shape, shape);
          if (!sameShape &&
              !tfc.util.arraysEqual(tensor.squeeze().shape, squeezedShape)) {
            throw new Error('the input tensors shape does not match');
          }
          return sameShape ? tensor : tensor.reshape(shape);
        });
        return [tfc.stack(mapped, axis)];
      });
    }
    case 'unstack': {
      return tfc.tidy(() => {
        const axis = getParamValue('axis', node, tensorMap, context) as number;
        const tensor =
            getParamValue('tensor', node, tensorMap, context) as tfc.Tensor;
        return tfc.unstack(tensor, axis);
      });
    }
    case 'tile': {
      const reps = getParamValue('reps', node, tensorMap, context) as number[];
      return [tfc.tile(
          getParamValue('x', node, tensorMap, context) as tfc.Tensor, reps)];
    }
    case 'split': {
      const axis = getParamValue('axis', node, tensorMap, context) as number;
      const numOrSizeSplits =
          getParamValue('numOrSizeSplits', node, tensorMap, context) as number;
      return tfc.split(
          getParamValue('x', node, tensorMap, context) as tfc.Tensor,
          numOrSizeSplits, axis);
    }
    default:
      throw TypeError(`Node type ${node.op} is not implemented`);
  }
};