TensorFlow转Pytorch网络

前言

我在github搜了一个mnist分类代码，基于TensorFlow框架的，效果挺不错，但是我只想用pytorch，所以打算转成pytorch代码。

过程

TensorFlow网络：

def model_net(num):

    model = Sequential()

    model.add(Conv2D(filters=16, kernel_size=(5, 5), padding='Same',
                     activation='relu', input_shape=(28, 28, 3)))
    model.add(BatchNormalization())


    model.add(Conv2D(filters=16, kernel_size=(3, 3), padding='Same',
                     activation='relu'))
    model.add(BatchNormalization())
    model.add(Conv2D(filters=32, kernel_size=(3, 3), padding='Same',
                     activation='relu'))
    model.add(BatchNormalization())
    model.add(MaxPool2D(pool_size=(2, 2), strides=(2, 2)))


    model.add(Flatten())
    model.add(Dense(256, activation="relu"))
    model.add(Dropout(0.5))
    model.add(Dense(num, activation="softmax"))
    return model

step1：输出model的结构和参数

class_num=5

model = model_net(class_num)
model.summary()
from tensorflow.keras.utils import plot_model
plot_model(model, to_file='DResLayer_model.png', show_shapes=True)  
print(model)

结构如下： 

 step2：构造pytorch网络并输出结构：

模拟cnn构造一下：

class Mnist(nn.Module):
  def __init__(self):
    super(Mnist,self).__init__()

    self.conv1 = nn.Sequential(
          # 输入3，输出 16，核大小为5，步长1，填充边缘1，填充方式 0填充
          nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, 
                    stride=1, padding=2, bias=False),
          nn.ReLU(),
          # num_features取 通道数
          nn.BatchNorm2d(num_features=16,affine=False)
      )
    self.conv2 = nn.Sequential(
        nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3,
                  stride=1, padding=1, bias=False),
        nn.ReLU(),
        nn.BatchNorm2d(num_features=16,affine=False)
    )
    self.conv3 = nn.Sequential(
        nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3,
                  stride=1, padding=1, bias=False)
        , nn.ReLU(),
        nn.BatchNorm2d(num_features=32,affine=False),
        nn.MaxPool2d(kernel_size=2, stride=(2, 2)),
        nn.Flatten()
    )
    self.fc = nn.Sequential(
        nn.Linear(in_features=32 * 14 * 14, out_features=256),
        nn.ReLU(),
        nn.Dropout2d(0.5),
        nn.Linear(in_features=256, out_features=5)

    )

  def forward(self, x):
      x = self.conv1(x)
      x = self.conv2(x)
      x = self.conv3(x)
      x = x.view(x.size(0), -1)
      x = self.fc(x)
      return x

输出结构需要先 pip install torch-summary

model=Mnist()
print(model)

from torchsummary import summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
t = model.to(device)
summary(t, (3, 28, 28))

输出如下：

 step3：对比结构和参数

参数略有不同，原因嘛。。。。。。，但是发现运行结果还是差不多的。