如何为Tensorflow构建自定义数据集

licenses(["notice"])  # Apache 2.0
 
package(default_visibility = ["//visibility:public"])
 
cc_binary(
    name = "python/ops/_pcap_ops.so",
    srcs = [
        "kernels/pcap_input.cc",
        "ops/pcap_ops.cc",
    ],
    copts = [
        "-pthread",
        "-std=c++11",
        "-DNDEBUG",
    ],
    linkshared = 1,
    deps = [
        "//tensorflow_io/core:dataset_ops",
        "@local_config_tf//:libtensorflow_framework",
        "@local_config_tf//:tf_header_lib",
    ],
)

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
 
namespace tensorflow {
 
REGISTER_OP("PcapInput")
    .Input("source: string")
    .Output("handle: variant")
    .Attr("filters: list(string) = []")
    .Attr("columns: list(string) = []")
    .Attr("schema: string = ''")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
       c->set_output(0, c->MakeShape({c->UnknownDim()}));
       return Status::OK();
     });
 
REGISTER_OP("PcapDataset")
    .Input("input: T")
    .Input("batch: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("T: {string, variant} = DT_VARIANT")
    .SetIsStateful()
    .SetShapeFn([](shape_inference::InferenceContext* c) {
       c->set_output(0, c->MakeShape({}));
       return Status::OK();
     });
 
}  // namespace tensorflow

REGISTER_UNARY_VARIANT_DECODE_FUNCTION(PcapInput, "tensorflow::data::PcapInput");
 
REGISTER_KERNEL_BUILDER(Name("PcapInput").Device(DEVICE_CPU),
                        FileInputOp<PcapInput>);
REGISTER_KERNEL_BUILDER(Name("PcapDataset").Device(DEVICE_CPU),
                        FileInputDatasetOp<PcapInput, PcapInputStream>);

   // read packets from the file up to record_to_read or end of file
    while ((*record_read) < record_to_read) {
      int64 record_count = 0;
      double packet_timestamp;
      string packet_data_buffer;
      Status status = state.get()->ReadRecord(packet_timestamp, &packet_data_buffer, record_count);
      if (!(status.ok() || errors::IsOutOfRange(status))) {
        return status;
      }
      if (record_count > 0) {
        Tensor timestamp_tensor = (*out_tensors)[0];
        timestamp_tensor.flat<double>()(*record_read) = packet_timestamp;
        Tensor data_tensor = (*out_tensors)[1];
        data_tensor.flat<string>()(*record_read) = std::move(packet_data_buffer);
        (*record_read) += record_count;
      } else {
        // no more records available to read
        // record_count == 0
        break;
      }
    }

ReadRecord(packet_timestamp, &packet_data_buffer, record_count);

Tensor timestamp_tensor = (*out_tensors)[0];
timestamp_tensor.flat<double>()(*record_read) = packet_timestamp;
Tensor data_tensor = (*out_tensors)[1];
data_tensor.flat<string>()(*record_read) = std::move(packet_data_buffer);

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
 
from tensorflow.python.util.all_util import remove_undocumented
from tensorflow_io.pcap.python.ops.pcap_ops import PcapDataset
 
_allowed_symbols = [
    "PcapDataset",
]
 
remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)

"""PcapDataset"""
import tensorflow as tf
from tensorflow_io.core.python.ops import data_ops as data_ops
from tensorflow_io import _load_library
pcap_ops = _load_library('_pcap_ops.so')
 
 
class PcapDataset(data_ops.Dataset):
  """A pcap Dataset. Pcap is a popular file format for capturing network packets.
  """
 
  def __init__(self, filenames, batch=None):
    """Create a pcap Reader.
    Args:
      filenames: A `tf.string` tensor containing one or more filenames.
    """
    batch = 0 if batch is None else batch
    dtypes = [tf.float64, tf.string]
    shapes = [
        tf.TensorShape([]), tf.TensorShape([])] if batch == 0 else [
            tf.TensorShape([None]), tf.TensorShape([None])]
    super(PcapDataset, self).__init__(
        pcap_ops.pcap_dataset,
        pcap_ops.pcap_input(filenames),
        batch, dtypes, shapes)

from tensorflow_io import _load_library
pcap_ops = _load_library('_pcap_ops.so')

dtypes = [tf.float64, tf.string]

shapes = [
        tf.TensorShape([]), tf.TensorShape([])] if batch == 0 else [
            tf.TensorShape([None]), tf.TensorShape([None])]

 dataset = pcap_io.PcapDataset(url_filenames, batch=1)
 
  packets_total = 0
  for v in dataset:
    (packet_timestamp, packet_data) = v
    if packets_total == 0:
      assert packet_timestamp.numpy()[0] == 1084443427.311224 # we know this is the correct value in the test pcap file
      assert len(packet_data.numpy()[0]) == 62 # we know this is the correct packet data buffer length in the test pcap file
    packets_total += 1
  assert packets_total == 43 # we know this is the correct number of packets in the test pcap file

$ bazel build -s --verbose_failures //tensorflow_io/pcap/...
$ pytest tests/test_pcap_eager.py