alexanderlutsenko / nobuco Goto Github PK

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
TransformerEncoderLayer                  [1, 512, 128]             --
├─MultiheadAttention: 1-1                [1, 512, 128]             66,048
├─Dropout: 1-2                           [1, 512, 128]             --
├─LayerNorm: 1-3                         [1, 512, 128]             256
├─Linear: 1-4                            [1, 512, 256]             33,024
├─Dropout: 1-5                           [1, 512, 256]             --
├─Linear: 1-6                            [1, 512, 128]             32,896
├─Dropout: 1-7                           [1, 512, 128]             --
├─LayerNorm: 1-8                         [1, 512, 128]             256
==========================================================================================
Total params: 132,480
Trainable params: 132,480
Non-trainable params: 0
Total mult-adds (M): 0.07
==========================================================================================
Input size (MB): 0.26
Forward/backward pass size (MB): 2.62
Params size (MB): 0.27
Estimated Total Size (MB): 3.15
==========================================================================================

__________________________________________________________________________________________________
 Layer (type)                Output Shape                 Param #   Connected to                  
==================================================================================================
 input_1 (InputLayer)        [(1, 128, 512)]              0         []                            
                                                                                                  
 tf.compat.v1.transpose_2 (  (1, 512, 128)                0         ['input_1[0][0]']             
 TFOpLambda)                                                                                      
                                                                                                  
 multi_head_attention_2 (Mu  (1, 512, 128)                66048     ['tf.compat.v1.transpose_2[0][
 ltiHeadAttention)                                                  0]',                          
                                                                     'tf.compat.v1.transpose_2[0][
                                                                    0]',                          
                                                                     'tf.compat.v1.transpose_2[0][
                                                                    0]']                          
                                                                                                  
 dropout_1 (Dropout)         (1, 512, 128)                0         ['multi_head_attention_2[1][0]
                                                                    ']                            
                                                                                                  
 tf.compat.v1.transpose_3 (  (1, 128, 512)                0         ['dropout_1[0][0]']           
 TFOpLambda)                                                                                      
                                                                                                  
 weight_layer_4 (WeightLaye  (1, 512, 256)                131072    ['input_1[0][0]']             
 r)                                                                                               
                                                                                                  
 tf.__operators__.add (TFOp  (1, 128, 512)                0         ['input_1[0][0]',             
 Lambda)                                                             'tf.compat.v1.transpose_3[0][
                                                                    0]']                          
                                                                                                  
 dropout_2 (Dropout)         (1, 512, 256)                0         ['weight_layer_4[0][0]']      
                                                                                                  
 tf.compat.v1.transpose_4 (  (1, 512, 128)                0         ['tf.__operators__.add[0][0]']
 TFOpLambda)                                                                                      
                                                                                                  
 dense_1 (Dense)             (1, 512, 128)                32896     ['dropout_2[0][0]']           
                                                                                                  
 layer_normalization (Layer  (1, 512, 128)                256       ['tf.compat.v1.transpose_4[0][
 Normalization)                                                     0]']                          
                                                                                                  
 dropout_3 (Dropout)         (1, 512, 128)                0         ['dense_1[0][0]']             
                                                                                                  
 tf.__operators__.add_1 (TF  (1, 512, 128)                0         ['layer_normalization[0][0]', 
 OpLambda)                                                           'dropout_3[0][0]']           
                                                                                                  
 layer_normalization_1 (Lay  (1, 512, 128)                256       ['tf.__operators__.add_1[0][0]
 erNormalization)                                                   ']                            
                                                                                                  
 tf.compat.v1.transpose_5 (  (1, 128, 512)                0         ['layer_normalization_1[0][0]'
 TFOpLambda)                                                        ]                             
                                                                                                  
 tf.identity (TFOpLambda)    (1, 128, 512)                0         ['tf.compat.v1.transpose_5[0][
                                                                    0]']                          
                                                                                                  
==================================================================================================
Total params: 230528 (900.50 KB)
Trainable params: 230528 (900.50 KB)
Non-trainable params: 0 (0.00 Byte)
__________________________________________________________________________________________________

  File "/home/titanx/hengck/opt/anaconda3.9/lib/python3.9/site-packages/nobuco/converters/node_converter.py", line 48, in decorator
    converter_result_func = converter_func(*args, **kwargs)
TypeError: converter_mean() got an unexpected keyword argument 'keepdims'

InvalidArgumentError: Exception encountered when calling layer 'depthwise_conv2d_1' (type DepthwiseConv2D).

{{function_node __wrapped__DepthwiseConv2dNative_device_/job:localhost/replica:0/task:0/device:CPU:0}} Current implementation only supports equal length strides in the row and column dimensions. [Op:DepthwiseConv2dNative] name: 

Call arguments received by layer 'depthwise_conv2d_1' (type DepthwiseConv2D):
  • inputs=tf.Tensor(shape=(4, 10, 10, 12), dtype=float32)

TypeError: '<' not supported between instances of 'tuple' and 'int'

@converter(torch.nn.functional.log_softmax, channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)
def converter_log_softmax(input, dim, dtype=None):
    def func(input, dim, dtype=None):
        # Adjust 'dim' if it's negative to handle PyTorch's negative indexing
        if dim < 0:
            dim += len(input.shape)

        # Apply TensorFlow's log_softmax
        # If dtype is specified, cast the input tensor to this dtype first
        if dtype is not None:
            input = tf.cast(input, dtype)
        return tf.nn.log_softmax(input, axis=dim)

    return func

/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py:55: RuntimeWarning: [<class 'lightglue.lightglue.TransformerLayer'>|LightGlue] conversion procedure might be incorrect: max. discrepancy for output #1 is 0.00010 (0.004%)
  warnings.warn(warn_string, category=RuntimeWarning)
/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py:55: RuntimeWarning: [<class 'lightglue.lightglue.TransformerLayer'>|LightGlue] conversion procedure might be incorrect: max. discrepancy for output #0 is 0.00012 (0.005%)
  warnings.warn(warn_string, category=RuntimeWarning)
/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py:55: RuntimeWarning: [<function log_softmax at 0x7bc0c800de10>|LightGlue->MatchAssignment] conversion procedure might be incorrect: max. discrepancy for output #0 is 38.75780 (103.477%)
  warnings.warn(warn_string, category=RuntimeWarning)
/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py:55: RuntimeWarning: [<class 'lightglue.lightglue.MatchAssignment'>|LightGlue] conversion procedure might be incorrect: max. discrepancy for output #0 is 38.75780 (43.686%)
  warnings.warn(warn_string, category=RuntimeWarning)

# Original implementation
# def sigmoid_log_double_softmax(
#     sim: torch.Tensor, z0: torch.Tensor, z1: torch.Tensor
# ) -> torch.Tensor:
#     """create the log assignment matrix from logits and similarity"""
#     b, m, n = sim.shape
#     certainties = F.logsigmoid(z0) + F.logsigmoid(z1).transpose(1, 2)
#     scores0 = F.log_softmax(sim, 2)
#     scores1 = F.log_softmax(sim.transpose(-1, -2).contiguous(), 2).transpose(-1, -2)
#     scores = sim.new_full((b, m + 1, n + 1), 0)
#     scores[:, :m, :n] = scores0 + scores1 + certainties
#     scores[:, :-1, -1] = F.logsigmoid(-z0.squeeze(-1))
#     scores[:, -1, :-1] = F.logsigmoid(-z1.squeeze(-1))
#     return scores

# My implementation with some modifications to eliminate slicing
def sigmoid_log_double_softmax(sim: torch.Tensor, z0: torch.Tensor, z1: torch.Tensor) -> torch.Tensor:
    """create the log assignment matrix from logits and similarity"""
    b, m, n = sim.shape

    # Calculate certainties and scores0, scores1 as before
    certainties = F.logsigmoid(z0) + F.logsigmoid(z1).transpose(1, 2)
    scores0 = F.log_softmax(sim, 2)
    scores1 = F.log_softmax(sim.transpose(-1, -2).contiguous(), 2).transpose(-1, -2)

    # Create scores tensor
    scores = sim.new_full((b, m + 1, n + 1), 0)

    # Merge the scores0, scores1, and certainties into scores without slice assignment
    scores_main = scores0 + scores1 + certainties
    scores[:, :m, :n] = scores_main

    # Compute the scores for the last column and row
    last_col_scores = F.logsigmoid(-z0.squeeze(-1)).unsqueeze(2)
    last_row_scores = F.logsigmoid(-z1.squeeze(-1)).unsqueeze(1)

    # Update last column and row in scores
    scores[:, :-1, -1:] = last_col_scores
    scores[:, -1:, :-1] = last_row_scores

    return scores
    ```

WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.linalg.matmul), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(513, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.linalg.matmul), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(513, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.reshape_3), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(2, 1, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.reshape_3), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(2, 1, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.reshape_2), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(2, 1, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.reshape_2), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(2, 1, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.identity_1), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(1, 27, 1) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.identity_1), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(1, 27, 1) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.math.multiply_51), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(1,) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.
WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.math.multiply_51), but are not present in its tracked objects:   <tf.Variable 'weight:0' shape=(1,) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.

import torch
import torch.nn as nn
import torch.nn.functional as F
from nobuco.convert.converter import pytorch_to_keras
import numpy as np

class MHSA(nn.Module):
    def __init__(self,
            embed_dim,
            out_dim,
            qk_dim,
            v_dim,
            num_head,
        ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_head  = num_head
        self.qk_dim = qk_dim
        self.v_dim  = v_dim

        self.q = nn.Linear(embed_dim, qk_dim*num_head)
        self.k = nn.Linear(embed_dim, qk_dim*num_head)
        self.v = nn.Linear(embed_dim, v_dim*num_head)
        
        self.out = nn.Linear(v_dim*num_head, out_dim)
        self.scale = 1/(qk_dim**0.5)

    def forward(self, x):
        L,dim = x.shape
        num_head = self.num_head
        qk_dim = self.qk_dim
        v_dim = self.v_dim
        
        q = self.q(x)
        k = self.k(x)
        v = self.v(x)
        q = q.reshape(L, num_head, qk_dim).permute(1,0,2).contiguous()
        k = k.reshape(L, num_head, qk_dim).permute(1,2,0).contiguous()
        v = v.reshape(L, num_head, v_dim ).permute(1,0,2).contiguous()

        dot = q *self.scale @ k  # H L L
        attn = F.softmax(dot, -1)    # L L

        v = torch.matmul(attn, v)  # L H dim
        v = v.permute(1,0,2).reshape(L, v_dim*num_head).contiguous()
        out = self.out(v)

        return out
if __name__ == "__main__":
    emb_dim = 128
    out_dim = 200
    num_heads = 4
    qk_dim = emb_dim // num_heads
    L = 33
    model = MHSA(embed_dim=emb_dim, out_dim=out_dim, qk_dim=qk_dim, v_dim=qk_dim, num_head=num_heads)
    keras_model = pytorch_to_keras(
        model, args=[torch.rand(L, emb_dim)],
    )
    inp = np.random.rand(33, 128)
    print(keras_model(inp)) ## runs fine

    inp = np.random.rand(15, 128)
    keras_model(inp) ## fails

InvalidArgumentError: Exception encountered when calling layer 'tf.reshape' (type TFOpLambda).

{{function_node __wrapped__Reshape_device_/job:localhost/replica:0/task:0/device:GPU:0}} Input to reshape is a tensor with 1920 values, but the requested shape has 4224 [Op:Reshape]

Call arguments received by layer 'tf.reshape' (type TFOpLambda):
  • tensor=tf.Tensor(shape=(15, 128), dtype=float32)
  • shape=('33', '4', '32')
  • name=None

slime@slime:~$ pip freeze | grep -E "torch=|nobuco"
nobuco==0.1.1
torch==2.1.0.dev20230317+cu118

ValueError: Layer depthwise_conv1d weight shape (3, 32, 1) is not compatible with provided weight shape (3, 1, 32).

import torch.nn as nn
import torch

from nobuco.convert.converter import pytorch_to_keras
from nobuco.commons import ChannelOrder

class SimpleDepthwiseConvExample(nn.Module):
    def __init__(self, in_channels=32, out_channels=32, kernel_size=3, stride=1, n_groups=32):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, kernel_size//2, groups=n_groups)
    
    def forward(self, x):
        return self.conv(x)

if __name__ == '__main__':
    bs = 1
    n_channels = 32
    n_groups = n_channels
    # n_groups = 1
    features = 144
    inputs = [
        torch.rand(bs, n_channels, features)
    ]
    model = SimpleDepthwiseConvExample(
        in_channels=n_channels,
        out_channels=n_channels,
        n_groups=n_groups
    )
    keras_model = pytorch_to_keras(
        model, inputs, inputs_channel_order=ChannelOrder.PYTORCH
    )

import torch.nn as nn
import torch

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2


import tensorflow as tf
import nobuco
from nobuco.commons import ChannelOrder, ChannelOrderingStrategy

class ExampleModel(nn.Module):
    def __init__(self, 
                **kwargs):
        
        super(ExampleModel, self).__init__()
        self.layer1 = nn.Conv2d(16, 16, (3,3), (1,1), (0,0), (1,1), 16)
        self.layer2 = nn.ReLU()


    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x


model = ExampleModel()

# Put model in inference mode
model.eval()

x = torch.randn(1, 16, 113, 113, requires_grad=False)

keras_model = nobuco.pytorch_to_keras(
    model,
    args=[x], kwargs=None)

# Assuming 'model' is your Keras model
full_model = tf.function(lambda x: keras_model(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))

# Convert Keras model to frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

# Print the input and output tensors
print("Frozen model inputs: ", frozen_func.inputs)
print("Frozen model outputs: ", frozen_func.outputs)

# Save frozen graph to disk
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                logdir='.',
                name='ExampleModel.pb',
                as_text=False)

@nobuco.converter(nn.Conv2d)
def converter_Conv2d(self, input: Tensor):
    weight = self.weight
    bias = self.bias
    groups = self.groups
    padding = self.padding
    stride = self.stride
    dilation = self.dilation
    
    

    out_filters, in_filters, kh, kw = weight.shape

    weights = weight.cpu().detach().numpy()

    if groups == 1:
        weights = tf.transpose(weights, (2, 3, 1, 0))
    else:
        weights = tf.transpose(weights, (2, 3, 0, 1))

    if bias is not None:
        biases = bias.cpu().detach().numpy()
        params = [weights, biases]
        use_bias = True
    else:
        params = [weights]
        use_bias = False

    if isinstance(dilation, numbers.Number):
        dilation = (dilation, dilation)

    if isinstance(padding, numbers.Number):
        padding = (padding, padding)

    pad_str = 'valid'
    pad_layer = None

    if padding == 'same':
        pad_str = 'same'
    elif padding != (0, 0):
        pad_layer = keras.layers.ZeroPadding2D(padding)

    if groups == 1:
        conv = keras.layers.Conv2D(filters=out_filters,
                                kernel_size=(kh, kw),
                                strides=stride,
                                padding=pad_str,
                                dilation_rate=dilation,
                                groups=groups,
                                use_bias=use_bias,
                                weights=params
                                )
    else:
        conv = keras.layers.DepthwiseConv2D(
                    kernel_size=(kh, kw),
                    strides=stride,
                    padding=pad_str,
                    use_bias=use_bias,
                    activation=None,
                    depth_multiplier=1,
                    weights=params,
                    dilation_rate=dilation,
                )

    def func(input):
        if pad_layer is not None:
            input = pad_layer(input)
        output = conv(input)
        return output
    return func

def complex_initializer(base_initializer):
    f = base_initializer()

    def initializer(*args, dtype=tf.complex64, **kwargs):
        real = f(*args, **kwargs)
        imag = f(*args, **kwargs)
        return tf.complex(real, imag)

    return initializer
    
initializer=complex_initializer(tf.random_normal_initializer))

Traceback (most recent call last):
  File "/p/i/proj/waternet_tf.py", line 5, in <module>
    import nobuco
  File "/p/i/proj/.venv/lib/python3.12/site-packages/nobuco/__init__.py", line 1, in <module>
    from nobuco.converters.channel_ordering import t_pytorch2keras, t_keras2pytorch
  File "/p/i/proj/.venv/lib/python3.12/site-packages/nobuco/converters/channel_ordering.py", line 6, in <module>
    from nobuco.commons import ChannelOrder, TF_TENSOR_CLASSES
  File "/p/i/proj/.venv/lib/python3.12/site-packages/nobuco/commons.py", line 3, in <module>
    from keras.src.engine.keras_tensor import KerasTensor
ModuleNotFoundError: No module named 'keras.src.engine'

nobuco                       0.11.5
keras                        2.15.0
torch                        2.2.0+cu118
tensorflow                   2.15.0.post1

@nobuco.converter(torch.nn.functional.log_softmax, channel_ordering_strategy=nobuco.ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)
def log_softmax(input: torch.Tensor, dim: int = -1):
    return lambda input, dim=dim: tf.nn.log_softmax(input, axis=dim)

class OutBlock(nn.Module):
    def __init__(self):
        super(OutBlock, self).__init__()
        self.conv = nn.Conv2d(256, 1, kernel_size=1)  # value head
        self.bn = nn.BatchNorm2d(1)
        self.fc1 = nn.Linear(10 * 10, 128)
        self.fc2 = nn.Linear(128, 1)

        self.conv1 = nn.Conv2d(256, 128, kernel_size=1)  # policy head
        self.bn1 = nn.BatchNorm2d(128)
        self.logsoftmax = nn.LogSoftmax(dim=1)
        self.fc = nn.Linear(10 * 10 * 128, 10 * 10 * 36)

    def forward(self, s):
        v = F.relu(self.bn(self.conv(s)))  # value head
        v = v.view(-1, 10 * 10)  # batch_size X channel X height X width
        v = F.relu(self.fc1(v))
        v = torch.tanh(self.fc2(v))

        p = F.relu(self.bn1(self.conv1(s)))  # policy head
        p = p.view(-1, 10 * 10 * 128)
        p = self.fc(p)
        p = self.logsoftmax(p).exp()
        return p, v