Не могу создать CNN с PyTorch

Question

Я раньше не использовал PyTorch, и я хотел создать CNN, который является адаптацией VGG19, который я создал в TensorFlow 2.0 с Keras.

Это код моего CNN в PyTorch.

class VGG19(nn.Module):


  def __init__(self, num_clases=43):
    super(VGG19, self).__init__()


    self.layer1 = nn.Sequential(
      nn.Conv2d(32, 64, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(64),
      # Repetimos lo mismo
      nn.Conv2d(64, 64, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(64),
      nn.MaxPool2d(kernel_size=2, stride=2))


    self.layer2 = nn.Sequential(
      nn.Conv2d(64, 128, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(128),

      nn.Conv2d(128, 128, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(128))

    self.layer3 = nn.Sequential(
      nn.Conv2d(128, 256, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(256),
      nn.Conv2d(256, 256, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(256),
      nn.Conv2d(256, 256, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(256),
      nn.Conv2d(256, 256, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(256),
      nn.MaxPool2d(kernel_size=2, stride=2))

    self.layer4 = nn.Sequential(
      nn.Conv2d(256, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512)) 

    self.layer5 = nn.Sequential(
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.Conv2d(512, 512, kernel_size=3, padding=2, stride=1),
      nn.ReLU(),
      nn.BatchNorm2d(512),
      nn.MaxPool2d(kernel_size=2, stride=2))

    self.fc1 = nn.Sequential(
      nn.Linear(512 * 15 * 11, 4096),
      nn.ReLU(),
      nn.BatchNorm2d(4096),
      nn.Dropout(0.5))

    self.fc2 = nn.Sequential(
      nn.Linear(4096, 4096),
      nn.ReLU(),
      nn.BatchNorm2d(4096),
      nn.Dropout(0.5))

    self.final = nn.Sequential(
      nn.Linear(4096, num_clases),
      nn.Softmax(num_clases))


  def forward(self, x) :
    out = self.layer1(x)
    out = self.layer2(out)
    out = self.layer3(out)
    out = self.layer4(out)
    out = self.layer5(out)
    print(out.shape)
    out = out.view(out.size(0), -1)
    out = self.fc1(out)
    out = self.fc2(out)
    out = self.final(out)
    return out

Когда я пытаюсь обучить его, я получаю эту ошибку.

total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = vgg19(images)
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

ValueError                                Traceback (most recent call last)
<ipython-input-39-7cc3d2f9bfc9> in <module>()
      6 
      7         # Forward pass
----> 8         outputs = vgg19(images)
      9         loss = criterion(outputs, labels)
     10 

6 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/batchnorm.py in _check_input_dim(self, input)
    249         if input.dim() != 4:
    250             raise ValueError('expected 4D input (got {}D input)'
--> 251                              .format(input.dim()))
    252 
    253 

ValueError: expected 4D input (got 2D input)

С Keras это действительно легко соединить сверточный слой с полностью связанным с model.add (Flatten ()) y model.add (Dense (4096)). Здесь я пытаюсь сделать то же самое, но это не работает.

Кто-нибудь может мне помочь, пожалуйста?