Потеря Pytorch свертки 0,0 с начала

Question

Я строю убедительную сеть, которая классифицирует собак и кошек. Архитектура довольно проста. 2 Conv (с пакетной нормой, leakyReLU, Maxpooling) до 1 фк. Размер входного изображения изменен до 64. Размер хороший. Проблема в том, что проигрыш равен 0,0 с самого начала. Я понятия не имею, в чем причина. Я не мог найти ответ. Я написал каждую деталь, которая может быть важной. Если вам нужно что-то еще, пожалуйста, скажите мне, я буду редактировать.

main.py

import torch
import torch.nn as nn
from torchvision import transforms, datasets
import PIL
import matplotlib.pyplot as plt
from Dataset import Dataset
from Network import Network
# Added to avoid     torch._C._cuda_init() \n RuntimeError: CUDA error: unknown error
torch.cuda.current_device()
# Hyper Parameters
batch_size = 1
img_size = 64
learning_rate = 0.001
num_epoch = 1

# Directories
trainDir = "D:/Programming/python/Deep learning/datasets/dogs-vs-cats/train"
testDir = "D:/Programming/python/Deep learning/datasets/dogs-vs-cats/test1"

print("Initializing...")

# Device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Augmentation
transforms = transforms.Compose([
    transforms.Resize((img_size, img_size)),
    transforms.ColorJitter(hue=.05, saturation=.05),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(20, resample=PIL.Image.BILINEAR) ,
    transforms.ToTensor()
])

trainset = datasets.ImageFolder(root=trainDir, transform=transforms)
testset = datasets.ImageFolder(root=testDir, transform=transforms)
train_loader = torch.utils.data.DataLoader(
    trainset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    testset, batch_size=batch_size, shuffle=False)  # test set will not be shuffled

model = Network(img_size,2).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

total_step = len(train_loader)
print("Tranining started")
for epoch in range(num_epoch):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        # forward propagate
        outputs = model(images)
        loss = criterion(outputs, labels)

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

        if (i+1) % 100 == 0:
            print(
                "Epoch [{}/{}], Step[{}/{}], Loss: {}".format(
                    epoch+1, num_epoch, i+1, total_step, loss.item()
                )
            )
print("Tranining complete, validation started")

with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        print('Test Accuracy: {} %'.format(100 * correct / total))

#
torch.save(model.state_dict(), "model.ckpy")

Network.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

"""
    Input size for conv
    l = number of input feature maps
    k = number of output feature maps
    n, m = width and height of kernel
    total parameter = (n*m*l+1)*k
"""

class Network(nn.Module):
    def __init__(self, input_size, num_class):
        super(Network, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        ) # output size = (128, 128, 16)
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        ) # output size = (64, 64, 32)

        self.fc1 = nn.Linear(
            int((input_size/4)**2*32), num_class
        )

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

выход

Epoch [1/1], Step[5800/25000], Loss: 0.0
Epoch [1/1], Step[5900/25000], Loss: 0.0
Epoch [1/1], Step[6000/25000], Loss: 0.0
Epoch [1/1], Step[6100/25000], Loss: 0.0
Epoch [1/1], Step[6200/25000], Loss: 0.0
Epoch [1/1], Step[6300/25000], Loss: 0.0
Epoch [1/1], Step[6400/25000], Loss: 0.0
Epoch [1/1], Step[6500/25000], Loss: 0.0

Результат после каждого слоя

outputs of conv1,2
         [[ 3.0135e-01,  3.5849e-01,  4.7758e-01,  ...,  3.9759e-01,
            3.7988e-01,  9.7870e-01],
          [ 4.3010e-01,  6.0753e-03,  4.5642e-01,  ..., -8.5486e-04,
            4.4537e-02,  2.9074e-01],
          [ 3.8567e-01,  7.8431e-02,  2.3859e-01,  ..., -3.0013e-03,
           -5.5821e-03,  1.2284e-01],
          ...,
          [ 3.9181e-01,  3.9093e-01,  1.2053e-01,  ..., -4.7156e-03,
            5.6266e-01,  7.7017e-01],
outputs of fc1
          [[-0.0772,  0.2166]]

Answer 1

loss = criterion(output, target.view(-1))  # Flatten target

попробуйте это.

не могли бы вы удалить эти две строки?images = images.to (устройство) label = tags.to (устройство)

self.conv1 и 2 должны быть отправлены в cuda: self.conv1(2).cuda()