Yup that’s how it should work I believe!
The custom Dataset class looks like this:
class PD_Dataset(Dataset):
def __init__(self, csv_file, root_dir, transform=None):
self.pd_features = pd.read_csv(csv_file)
self.root_dir = root_dir
self.transform = transform
self.classes = ("PD", "Non-PD")
def __len__(self):
return len(self.pd_features)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = to_list()
img_name = os.path.join("",self.pd_features.iloc[idx, 0])
image = Image.open(img_name)
features = self.pd_features.iloc[idx, 1]
labels = self.pd_features.iloc[idx, 2]
if self.transform:
image = self.transform(image)
return image, features, labels
Now one thing I have to figure out is changing the features variable here so that its shape follows image variable’s shape. features variable looks like [(x1, x2, x3)]
, and I was told that I need to pad a tensor with the same values so that shape becomes the shape of image.
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.cnn = models.resnet18(pretrained=True)
self.cnn.fc = nn.Sequential(
nn.Linear(self.cnn.fc.in_features, 256),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(256, 2),
nn.LogSoftmax(dim=1))
def forward(self, image, data):
x1 = self.cnn(image)
x2 = data
x = torch.cat((x1, x2), dim=1)
x = self.cnn.fc(x)
return x
model = MyModel()
Now for the training part:
for epoch in range(n_epochs):
train_loss = 0.0
valid_loss = 0.0
train_acc = 0
valid_acc = 0
model.train()
start = timer()
for ii, (data, target) in enumerate(train_loader):
if train_on_gpu:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
output = model(data)
target_tensor = torch.tensor(target, dtype=torch.long, device=torch.device('cuda'))
loss = criterion(output, target_tensor)
loss.backward()
optimizer.step()
train_loss += loss.item() * data.size(0)
_, pred = torch.max(output, dim=1)
correct_tensor = pred.eq(target.data.view_as(pred))
accuracy = torch.mean(correct_tensor.type(torch.FloatTensor))
train_acc += accuracy.item() * data.size(0)
Since my Dataset now will return in the format of image, features, data
, do you I have to enumerate the train_loader
a different way? And manually handle it? The reason I ask is, I am not seeing where the forward
function’s inputs are getting taken care of.
Would appreciate some help on this part!