I am trying to feed my own image dataset into PyTorch LSTM implementation that works fine with MNIST dataset. But getting " RuntimeError: Assertion `THIndexTensor_(size)(target, 0) == batch_size’ failed. at /Users/soumith/code/builder/wheel/pytorch-src/torch/lib/THNN/generic/ClassNLLCriterion.c:79" this error message.
import torch
import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as npfrom torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as dataTRANSFORM_IMG = transforms.Compose([
transforms.Resize(28,28),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225] )
])BATCH_SIZE = 4
TRAIN_DATA_PATH = “./train_set/”
TEST_DATA_PATH = “./test_set/”
train_data = torchvision.datasets.ImageFolder(root=TRAIN_DATA_PATH, transform=TRANSFORM_IMG)
train_data_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)test_data = torchvision.datasets.ImageFolder(root=TEST_DATA_PATH, transform=TRANSFORM_IMG)
test_data_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)labels = [‘low’,‘medium’,‘high’]
Hyper Parameters
sequence_length = 28
input_size = 28
hidden_size = 128
num_layers = 2
num_classes = 3num_epochs = 2
learning_rate = 0.01RNN Model (Many-to-One)
class RNN(nn.Module):
def init(self, input_size, hidden_size, num_layers, num_classes):
super(RNN, self).init()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
self.fc = nn.Linear(hidden_size, num_classes)def forward(self, x): # Set initial states h0 = Variable(torch.zeros(self.num_layers, x.size(0), self.hidden_size)) c0 = Variable(torch.zeros(self.num_layers, x.size(0), self.hidden_size)) # Forward propagate RNN out, _ = self.lstm(x, (h0, c0)) # Decode hidden state of last time step out = self.fc(out[:, -1, :]) return out
rnn = RNN(input_size, hidden_size, num_layers, num_classes)
Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)Train the Model
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(trainloader):
images = Variable(images.view(-1, sequence_length, input_size))
labels = Variable(labels)# Forward + Backward + Optimize optimizer.zero_grad() outputs = rnn(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() if (i+1) % 100 == 0: print ('Epoch [%d/%d], Step [%d/%d], Loss: %.4f' %(epoch+1, num_epochs, i+1, len(train_dataset)//batch_size, loss.data[0]))