玩转pytorch和tensorflow之（2）——conv2d

        二维卷积操作常用于特征提取，输入通常是一个四维的张量。tensorflow和pytorch的二维卷积的差异主要有两个：

        1）对于tensorflow来说，是BHWC的形状，通道数是最后一维；而对pytorch而言，输入是BCHW形状，通道数在第二维。

        2）对于tensorflow来说，weight权重的形状是kernel,  input_chan, output_chan

        例如对于(2,3)的卷积核,输入4通道输出8通道，则weight权重的形状是(2,3,4,8)

        而同样的配置下，pytorch的weight权重的形状是(8,4,2,3)

        bias的形状是一样的。

        至于padding，tensorflow有'valid'和'same'，pytorch直接是数字量。

        tensorflow的'valid'表示不padding，对应pytorch是(0,0),'same'表示输出形状和输入保持一致，pytorch需要根据输出输入的维度变化公式计算padding的数值。

        例子中 tensorflow padding是用'valid'模式。对应pytorch:padding=(0,0)

1)tensorflow转pytorch

weight参数要做一个转置transpose(3,2,0,1)，bias是一样的。

#tensorflow 参数转 pytorch
layer_th.weight.data = torch.tensor(weights[0].transpose(3,2,0,1))
layer_th.bias.data = torch.tensor(weights[1])

pytorch输入也要做一个转置transpose(0,3,1,2)，输出再转置回来，保持和tensorflow是一样的。

inputs_th = torch.from_numpy(inputs).permute(0,3,1,2)

with torch.no_grad():
    out = layer_th(inputs_th)
    outputs_th = out.permute(0,2,3,1).numpy()

下面是完整的代码：

#1)tensorflow 转为 pytorch
B,C,H,W=1,4,5,6
input_chan=C
output_chan=8

inputs = np.random.rand(B,H,W,C).astype("float32")

#tensorflow Conv2D 模型

#先用模型处理一次，这样才能得到模型的参数
inputs_tf = tf.convert_to_tensor(inputs)
layer_tf = tf.keras.layers.Conv2D(output_chan, (2, 3), (1, 2), padding="valid")
outputs_tf = layer_tf(inputs_tf, training=False)

#把模型的参数随机化
weights=layer_tf.get_weights()
weights[0]=np.random.rand(weights[0].shape[0],weights[0].shape[1],weights[0].shape[2],weights[0].shape[3]).astype("float32")
weights[1]=np.random.rand(weights[1].shape[0]).astype("float32")
layer_tf.set_weights(weights)
outputs_tf = layer_tf(inputs_tf, training=False)
outputs_tf = outputs_tf.numpy()

inputs_th = torch.from_numpy(inputs).permute(0,3,1,2)


#pytorch Conv2d 模型
layer_th = torch.nn.Conv2d(input_chan,output_chan,kernel_size=(2,3),stride=(1,2),padding=(0,0))

#tensorflow 参数转 pytorch
layer_th.weight.data = torch.tensor(weights[0].transpose(3,2,0,1))#.transpose(2,3,1,0).transpose(3,0,1,2)
layer_th.bias.data = torch.tensor(weights[1])


layer_th.eval()
with torch.no_grad():
    out = layer_th(inputs_th)
    outputs_th = out.permute(0,2,3,1).numpy()
    
print('tensorflow=>pytorch(conv2d):',np.allclose(outputs_th, outputs_tf, atol=1e-5))

 输出打印：tensorflow=>pytorch(conv2d): True

True表示转换成功。

2)pytorch转tensorflow

 weight参数要做一个转置permute(2,3,1,0)，bias是一样的。

#pytorch参数转到tensorflow
with torch.no_grad():
    keras_format_weights = [
                        layer_th.weight.permute(2,3,1,0).numpy(),
                        layer_th.bias.numpy()    ]

#2)pytorch 转为 tensorflow
#pytorch Conv2d 模型
layer_th = torch.nn.Conv2d(input_chan,output_chan,kernel_size=(2,3),stride=(1,2),padding=(0,0))
layer_th.eval()
with torch.no_grad():
    out = layer_th(inputs_th)
    outputs_th = out.permute(0,2,3,1).numpy()


#pytorch参数转到tensorflow
with torch.no_grad():
    keras_format_weights = [
                        layer_th.weight.permute(2,3,1,0).numpy(),
                        layer_th.bias.numpy()    ]

#tensorflow Conv2D 模型
model = keras.Sequential()
model.add(tf.keras.Input(shape=(H,W,C)))
model.add(tf.keras.layers.Conv2D(output_chan, (2, 3), (1, 2), padding="valid"))
model.layers[0].set_weights(keras_format_weights)
outputs_tf = model.predict(inputs_tf)

print('pytorch=>tensorflow(conv2d):',np.allclose(outputs_th, outputs_tf, atol=1e-5))                        

  输出打印：pytorch=>tensorflow(conv2d): True

True表示转换成功。

        至此，tensorflow和pytorch的二维卷积左右互搏之术已经修炼成功。至于是用左（tensorflow）还是右（pytorch），完全取决于你自己了。

上一篇《玩转pytorch和tensorflow之(1)——batchnorm-CSDN博客》

下一篇《玩转pytorch和tensorflow之(3)——layernorm-CSDN博客》