Function _pickle_save

python/paddle/framework/io.py:423–506

(obj, f, protocol)

def _pickle_save(obj, f, protocol):
    # TODO(weixin):add support for BytesIO.
    if not isinstance(protocol, int):
        raise ValueError(
            f"The 'protocol' MUST be `int`, but received {type(protocol)}"
        )

    if protocol < 2 or protocol > 4:
        raise ValueError(
            f"Expected 1<'protocol'<5, but received protocol={protocol}"
        )

    def reduce_varbase(self):
        if (
            self.is_dense()
            and self.place.is_custom_place()
            and core.is_compiled_with_custom_device('npu')
        ):
            data = np.array(paddle._C_ops.npu_identity(self, -1).cpu())
        else:
            data = np.array(self.cpu())
        name = self.name

        return (tuple, ((name, data),))

    def reduce_DenseTensor(self):
        p = core.Place()
        p.set_place(paddle.CPUPlace())
        if self._place().is_custom_place():
            data = np.array(paddle._C_ops.npu_identity(self, -1)._copy(p))
        else:
            data = np.array(self._copy(p))

        return (_reconstruct_dense_tensor_data, (data,))

    def reduce_Layer(self):
        raise ValueError(
            "paddle do not support saving `paddle.nn.Layer` object."
        )

    dispatch_table_layer = {}

    def create_layer_dispatch_table(layer):
        dispatch_table_layer[layer.__class__] = reduce_Layer
        return layer

    _parse_every_object(
        obj,
        lambda v: isinstance(v, paddle.nn.Layer),
        create_layer_dispatch_table,
    )

    def add_dispatch_table():
        # This is not a good method, because the pickle module has been modified.
        pickle.dispatch_table[core.eager.Tensor] = reduce_varbase
        pickle.dispatch_table[EagerParamBase] = reduce_varbase
        pickle.dispatch_table[core.DenseTensor] = reduce_DenseTensor
        pickle.dispatch_table.update(dispatch_table_layer)