Thank @albanD
I have bug from layer 8 to layer 9. Layer 9: fully connected with 4096 neurons.

error.jpg1612×472 214 KB

my code
    class Model(nn.Module):
    def __init__(self, num_classes=5):
        super(Model, self).__init__()
        
        #[256,256,3] (conv1)--> [63,63,96] (pool2)--> [31,31,96]
        self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2)
        self.lu1 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
        
        #[31,31,96] (conv3)--> [31,31,192] (pool4)-->[15,15,192]
        self.conv3 = nn.Conv2d(96, 192, kernel_size=5, padding=2, stride=1)
        self.lu3 = nn.ReLU()
        self.pool4 = nn.MaxPool2d(kernel_size=3, stride=2)
        
        #[15,15,192] (conv5) --> [15,15,288] 
        self.conv5 = nn.Conv2d(192, 288, kernel_size=3, stride=1, padding=1)
        self.lu5 = nn.ReLU()
        
        #[15,15,288] (conv6) --> [15,15,288]
        self.conv6 = nn.Conv2d(288, 288, kernel_size=3, stride=1, padding=1)
        self.lu6 = nn.ReLU()
        
        #[15,15,288] (conv7)--> [15,15,192] (pool8) --> [7,7,192]
        self.conv7 = nn.Conv2d(288, 192, kernel_size=3, stride=1, padding=1)
        self.lu7 = nn.ReLU()        
        self.pool8 = nn.MaxPool2d(kernel_size=3, stride=2)        
        
        # [7,7,192]
        self.fc9 = nn.Linear(in_features=192, out_features=4096)
        self.fc10 = nn.Linear(in_features=4096, out_features=4096)
        self.fc11 = nn.Linear(in_features=4096, out_features=5)
        
    def forward(self, input):
        output = self.conv1(input)
        output = self.lu1(output)
        output = self.pool2(output)
        
        output = self.conv3(output)
        output = self.lu3(output)
        output = self.pool4(output)
        
        output = self.conv5(output)
        output = self.lu5(output)
        
        output = self.conv6(output)
        output = self.lu6(output)
        
        output = self.conv7(output)
        output = self.lu7(output)
        output = self.pool8(output)
                
        output = self.fc9(output)
        output = self.fc10(output)
        output = self.fc11(output)        
        return output
        

epoch_nums = 60        
model = Model()

optimizer = Adam(model.parameters(), lr=5e-8, weight_decay=5e-1)     
loss_fn = nn.CrossEntropyLoss()

total_step = len(dataset_loader)

for epoch in range(epoch_nums):
    for i, (images, labels) in enumerate(dataset_loader):
        optimizer.zero_grad = 0
        outputs = model(images)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()