Package: tf-dataclass
tf-dataclassSupport python dataclass containers as input and output in callable TensorFlow 2 graph.InstallMake sure thattensorflow>=2.0.0ortensorflow-gpu>=2.0.0is installed.$pipinstalltf-dataclassWhyTensorFlow 2autograph functionsupports only nested structures of python tuples as inputs and output.
(Outputs can be also python dictionaries.)
This is inconvenient once we go beyond small hello world cases,
because we have to work with unstructured armfuls of tensors.
This small package is dedicated to fill this gap by letting@tf.functiondecorated functions to take and return pythonicdataclassinstancies.Examples of usage:1. Sequential featuresimporttensorflowastfimporttf_dataclass# Batch of sequential features of different length@tf_dataclass.dataclassclassSequential:feature:tf.Tensor# shape = [batch, length, channels],    dtype = tf.float32length:tf.Tensor# shape = [batch],                      dtype = tf.int32# Initialize a batch of two sequences of lengths 6 and 4input=Sequential(feature=tf.random.normal(shape=[2,6,3]),length=tf.constant([6,4],dtype=tf.int32),)# Define a convolution operator with a stride such that length -> length / stride@tf_dataclass.functiondefconvolution(input:Sequential,filters:tf.Tensor,stride:int)->Sequential:returnSequential(feature=tf.nn.conv1d(input.feature,filters,stride),length=tf.math.floordiv(input.length,stride),)# Output is an instance of Sequential with lengths 3 and 2 due to convolution stride = 2output=convolution(input=input,filters=tf.random.normal(shape=[1,3,7]),stride=2,)assertisinstance(output,Sequential)print(output.length)# -> tf.Tensor([3 2], shape=(2,), dtype=int32)2. Minibatch as a data transfer object:importtensorflowastfimporttf_dataclass@tf_dataclass.dataclassclassDataBatch:image:tf.Tensor# shape = [batch, height, width, channels], dtype = tf.flaot32label:tf.Tensor# shape = [batch],                          dtype = tf.int32image_file_path:tf.Tensor# shape = [batch],                          dtype = tf.stringdataset_name:tf.Tensor# shape = [batch],                          dtype = tf.string...@tf_dataclass.functiondeftrain_step(input:DataBatch)->None:...3. Containerized outputs:importtensorflowastfimporttf_dataclass@tf_dataclass.dataclassclassModelOutput:loss_value:tf.Tensor# shape = [batch],  dtype = tf.flaot32label:tf.Tensor# shape = [batch],  dtype = tf.int32prediction:tf.Tensor# shape = [batch],  dtype = tf.int32...@propertydefmean_loss(self)->tf.Tensor:# shape = [batch],  dtype = tf.float32returntf.reduce_mean(self.loss_value)@propertydefnum_true_predictions(self)->tf.Tensor:# shape = [batch],  dtype = tf.int32returntf.reduce_sum(tf.cast(self.label==self.prediction,dtype=tf.int32))@propertydefnum_false_predictions(self)->tf.Tensor:# shape = [batch],  dtype = tf.int32returntf.reduce_sum(tf.cast(self.label!=self.prediction,dtype=tf.int32))...@tf_dataclass.functiondefget_loss(...)->ModelOutput:...Such containers can be merged along datasets and workers.4. Easy tensorflow shape and dtype runtime verification:importtensorflowastfimporttf_dataclass@tf_dataclass.dataclassclassSequential:feature:tf.Tensor# shape = [batch, length, channels], dtype = tf.flaot32length:tf.Tensor# shape = [batch]                    dtype = tf.int32def__post_init__(self):# Verify featureassertself.feature.dtype==tf.float32assertlen(self.feature.shape)==3# Verify lengthassertself.length.dtype==tf.int32assertlen(self.length.shape)==1# Verify batch size# Works only in eager mode for better perfomanceassertself.feature.shape[0]==self.length.shape[0]@propertydefbatch_size(self)->tf.Tensor:# shape = [], dtype = tf.int32returntf.shape(self.feature)[0]Other features:Support hierarchical composition.Support inheritance including multiple one (for free from originaldataclass).Highliting, autocomplete, and refactoring from your IDE.UsageImportdataclassandfunctionfromtf_dataclassfromtf_dataclassesimportdataclass,functionIt is mandatory to use return type hints for the function decorated with@function.For example,fromtypingimportTuple@dataclassclassMyDataclass:...@functiondefmy_func(...)->Tuple[tf.Tensor,MyDataclass]:...returnsome_tensor,my_dataclass_instanceType hints for the arguments are optional but recommended.Positional arguments are not currently supported:For example, for@functiondefmy_graph_func(x:...,y:...)->...:...typemy_graph_func(x=x,y=y)but notmy_graph_func(x,y)Known ProblemsIDE autocomplete is currently not well-supported, for example, in PyCharm.
Solution: use importfromtypingimportTYPE_CHECKINGifTYPE_CHECKING:fromdataclassesimportdataclasselse:fromtf_dataclassimportdataclassin each*.pyfile wheredataclassis used.Under the roofDataclasses and their nested structures are simply converted into nested pythonic tuples and back.
This way we wrap given functions such that all inputs and outputs are nested tuples.
Then@tf.functionis applied. Afterward the graph function is wrapped bach to dataclass form.
Type hints are used in python runtime for the graph creation as temples to pack and unpack dataclass arguments.Future plansSupporttf.cond,tf.case,tf.switch_case,tf.while_loop,tf.Optional, andtf.data.Iterator.Support positional arguments.Conversion totf.neststructures.