PyTorch pretrained models not working

ppstain: <p>Hello, I’m trying to use PyTorch inbuilt models to categorise images, but the loss isn’t decreasing for cross entropy loss and is decreasing for mse loss. I’ve tried adding softmax and different loss functions but only mse loss works. This problem persists for other models like efficient net.<br>
Cross entropy loss<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd264dbb4b944fe9e8b6f4674b9e7d8fef03d199.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bd264dbb4b944fe9e8b6f4674b9e7d8fef03d199" title="Screen Shot 2022-10-26 at 11.36.55 am"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd264dbb4b944fe9e8b6f4674b9e7d8fef03d199_2_583x500.jpeg" alt="Screen Shot 2022-10-26 at 11.36.55 am" data-base62-sha1="qZij2qxw9FkM7kmPhOlDK25LGGJ" width="583" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd264dbb4b944fe9e8b6f4674b9e7d8fef03d199_2_583x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd264dbb4b944fe9e8b6f4674b9e7d8fef03d199_2_874x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd264dbb4b944fe9e8b6f4674b9e7d8fef03d199_2_1166x1000.jpeg 2x" data-dominant-color="E5EEF3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-10-26 at 11.36.55 am</span><span class="informations">1272×1090 82.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>This is the code</p>
<pre><code class="lang-auto">weights = AlexNet_Weights.DEFAULT
transform = weights.IMAGENET1K_V1.transforms()
model = alexnet(weights=None)
model.to(device)
data = '/Users/me/Downloads/archive/train'
dataset = datasets.ImageFolder(data, transform=transform)
model.classifier[6] = nn.Linear(4096,450)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True, pin_memory=True, num_workers=0)

images, labels = next(iter(dataloader))
loss_function = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),weight_decay=0.005, lr=0.005, momentum=0.9)
model.eval()

for epoch in range(50):
    data_x = []
    data_y = []
    images, labels = next(iter(dataloader))
    for i in range(32):
        data_x.append(np.array(images[i]))
        yes = torch.zeros(450)
        yes[labels[i].item()] = 1
        data_y.append(np.array(yes))
    print(epoch)
    data_x,data_y = np.array(data_x),np.array(data_y)
    data_x,data_y = torch.as_tensor(data_x,device=device),torch.as_tensor(data_y,device=device)
    output_y = model(data_x)
    loss = loss_function(output_y,data_y)
    losses.append(loss.item())
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    
</code></pre>

ptrblck: <p>It seems you are only using the first batch of your <code>DataLoader</code>:</p>
<pre><code class="lang-python">images, labels = next(iter(dataloader))
</code></pre>
<p>in each epoch so you might want to change it and then play around with some hyperparameters.</p>

ppstain: <p>I checked the images and labels by printing them every epoch and they were different, plus I used this same system for previous neural networks and it worked fine</p>