（conv1d）张量和jax产生了相同输入的不同输出

Question

我正在尝试使用Conv1D函数在JAX和TensorFlow上进行反式卷发。我为CON1D_transposed操作阅读了JAX和TensorFlow的文档，但它们的输出对相同的输入产生了不同的输出。

我找不到问题是什么。而且我不知道哪一个产生正确的结果。请帮我。

我的JAX实现（JAX代码）

x = np.asarray([[[1, 2, 3, 4, -5], [1, 2, 3, 4, 5]]], dtype=np.float32).transpose((0, 2, 1))
filters = np.array([[[1, 0, -1], [-1,  0,  1]], 
                    [[1, 1,  1], [-1, -1, -1]]], 
                    dtype=np.float32).transpose((2, 1, 0))

kernel_rot = np.rot90(np.rot90(filters))

print(f"x strides:  {x.strides}\nfilters strides: {kernel_rot.strides}\nx shape: {x.shape}\nfilters shape: {filters.shape}\nx: \n{x}\nfilters: \n{filters}\n")

dn1 = lax.conv_dimension_numbers(x.shape, filters.shape,('NWC', 'WIO', 'NWC'))
print(dn1)

res = lax.conv_general_dilated(x,kernel_rot,(1,),'SAME',(1,),(1,),dn1)     

res = np.asarray(res)
print(f"result strides: {res.strides}\nresult shape: {res.shape}\nresult: \n{res}\n")

我的TensorFlow实现（TensorFlow代码）

x = np.asarray([[[1, 2, 3, 4, -5], [1, 2, 3, 4, 5]]], dtype=np.float32).transpose((0, 2, 1))
filters = np.array([[[1, 0, -1], [-1,  0,  1]], 
                    [[1, 1,  1], [-1, -1, -1]]], 
                    dtype=np.float32).transpose((2, 1, 0))

print(f"x strides:  {x.strides}\nfilters strides: {filters.strides}\nx shape: {x.shape}\nfilters shape: {filters.shape}\nx: \n{x}\nfilters: \n{filters}\n")
    
res = tf.nn.conv1d_transpose(x, filters, output_shape = x.shape, strides = (1, 1, 1), padding = 'SAME', data_format='NWC', dilations=1)

res = np.asarray(res)
print(f"result strides: {res.strides}\nresult shape: {res.shape}\nresult: \n{res}\n")

JAX输出输出

result strides: (40, 8, 4)
result shape: (1, 5, 2)
result: 
[[[ 0.  0.]
  [ 0.  0.]
  [ 0.  0.]
  [10. 10.]
  [ 0. 10.]]]

从TensorFlow的

result strides: (40, 8, 4)
result shape: (1, 5, 2)
result: 
[[[  5.  -5.]
  [  8.  -8.]
  [ 11. -11.]
  [  4.  -4.]
  [  5.  -5.]]]

Answer 1

函数 [filter_width，output_channels，in_channels] 。如果过滤器在上面的片段中被转移以满足此形状，则jax返回正确的结果，而计算dn1参数应为woi> woi（<强> w idth - o utput_channels- i nput_channels）而不是wio> wio（ w idth- i nput_channels- o utput_channels）。之后：

result.strides = (40, 8, 4)
result.shape = (1, 5, 2)
result: 
[[[ -5.,   5.],
  [ -8.,   8.],
  [-11.,  11.],
  [ -4.,   4.],
  [ -5.,   5.]]]

结果与TensorFlow不同，但是JAX的内核被翻转了，因此实际上是预期的。