Pytorch Argmax形状问题

如何解决Pytorch Argmax形状问题

import sys,os
import cv2
import numpy as np
from tqdm import tqdm
REBUILD_DATA = True
import matplotlib.pyplot as plt
class ArtOrNot():
f = open('Art2','w')
f.write('oof')
f.close()
IMG_SIZE = 50
ART = (r'C:\Users\Kyel\Desktop\Python projects\ART')
NOTART = (r'C:\Users\Kyel\Desktop\Python projects\NOTART')
LABELS = {ART: 1,NOTART: 0}
training_data = []
artcount = 0
notartcount = 0


def make_training_data(self):
    for label in self.LABELS:
        print(label)
        for f in tqdm(os.listdir(label)):
            try:
                path = os.path.join(label,f)
                img = cv2.imread(path)
                img = cv2.resize(img,(self.IMG_SIZE,self.IMG_SIZE))
                #gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
                #(img) = (grey)
                self.training_data.append([np.array(img),np.eye(2)[self.LABELS[label]]])
                if label == self.ART:
                    self.artcount += 1
                    print(self.artcount)
                elif label == self.NOTART:
                    self.notartcount += 1
            except Exception as e:
                pass
    #print(training_data)
    np.random.shuffle(self.training_data)
    np.save("training_data.npy",self.training_data)
    print("Art:",self.artcount)
    print("NotART",self.notartcount)
    

if REBUILD_DATA:
    artornot = ArtOrNot()
    artornot.make_training_data()

training_data = np.load("training_data.npy",allow_pickle=True)
print(len(training_data))

import matplotlib.pyplot as plt
plt.imshow(training_data[1][0])
plt.show()
import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
    super().__init__()
    self.conv1 = nn.Conv2d(1,32,5)
    self.conv2 = nn.Conv2d(32,64,5)
    self.conv3 = nn.Conv2d(64,128,5)
    
    
    x = torch.randn(50,50).view(-1,1,50,50)
    self._to_linear = None
    self.convs(x)
    
    
    
    self.fc1 = nn.Linear(self._to_linear,512)
    self.fc2 = nn.Linear(512,2)
    
    
def convs(self,x):
    x = F.max_pool2d(F.relu(self.conv1(x)),(2,2))
    x = F.max_pool2d(F.relu(self.conv2(x)),2))
    x = F.max_pool2d(F.relu(self.conv3(x)),2))
    print(x[0].shape)
    
    if self._to_linear is None:
        self._to_linear = x[0].shape[0]*x[0].shape[1]*x[0].shape[2]
        
    return x

def forward(self,x):
    x = self.convs(x)
    x = x.view(-1,self._to_linear)
    c = F.relu(self.fc1(x))
    x = self.fc2(x)
    #return F.softmax(x,dim =1)

net = Net() 
print(net)

import torch.optim as optim

optimizer = optim.Adam(net.parameters(),lr=0.001)
#loss_function = torch.nn.functional.mse_loss()


X = torch.Tensor([i[0] for i in training_data]).view(-1,50)
X = X/255.0
y = torch.Tensor([i[1] for i in training_data])

VAL_PCT = 0.2
val_size = int(len(X)*VAL_PCT)
print(val_size)

train_X = X[:-val_size]
#print(train_X)
train_y = X[:-val_size]
#print(train_y)
test_X = X[-val_size:]
print(test_X)
test_y = y[-val_size:]
print(test_y)
#print(X)
print(len(X))
#test_X = 8

#print(len(train_X),print(test_X))
#print(len(train_X))
#print(len(train_X))
#print(test_X)

BATCH_SIZE = 4
EPOCHS = 10

for epoch in range(EPOCHS):
    for i in tqdm(range(0,len(train_X),BATCH_SIZE)): # from 0,to the len of x,stepping BATCH_SIZE at a time. [:50] ..for now just to dev
        #print(f"{i}:{i+BATCH_SIZE}")
        batch_X = train_X[i:i+BATCH_SIZE].view(-1,50)
        batch_y = train_y[i:i+BATCH_SIZE]

        input = torch.randn(3,5,requires_grad=True)
    
        target = torch.randn(3,5)
        output = loss(input,target)
        output.backward()
        optimizer.step()
        print(output)
    #print(f"Epoch: {epoch}. Loss: {loss}")
#print(loss)
#print(len(train_X))
#print(len(test_X))

问题领域

correct = 0
total = 0

with torch.no_grad():
    for i in tqdm(range(len(test_X))):
        real_class = torch.argmax(test_y[i])
    
        net_out = net(test_X[i].view(10,50))[0]  # returns a list,predicted_class = torch.argmax(net_out)

        if predicted_class == real_class:
            correct += 1
        total += 1
print("Accuracy: ",round(correct/total,3))

错误消息：RuntimeError：形状'[10，50，1，50]'对于大小为2500的输入无效

所以我知道从神经网络出来的数据大小是2500，当我更改参数时，又遇到另一个错误，说它期望4D张量，而我只给了3D张量？有帮助吗？