Sure!
Loading Imagenette:
def load_imagenette(batch_size):
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Imagenet Values
])
trainset = datasets.ImageFolder('data/imagenette2/train', transform=transform)
validset = datasets.ImageFolder('data/imagenette2/val', transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)
validloader = torch.utils.data.DataLoader(validset, batch_size=batch_size, shuffle=True)
return trainloader, validloader
I use VGG-16 pretrained=True, freeze all layers, except the output layer which got changed to have output = 10 classes.
Other infos:
batch_size = 64
criterion = CrossEntropyLoss
optimizer = SGD with lr = 0.01, momentum= 0.9
Training Function:
def train_model(model, trainloader, validloader, num_epochs, optimizer, criterion, device):
starting_time = time.time()
val_acc_history = []
best_model_wts = copy.deepcopy(model.state_dict())
best_train_acc = 0.0
best_val_acc = 0.0
for epoch in range(num_epochs):
print('Epoch: {}/{}'.format(epoch+1, num_epochs))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
model.train() # set model to training mode
dataloader = trainloader
else:
model.eval() # set model to eval mode
dataloader = validloader
running_loss = 0
running_correct = 0
for images, labels in dataloader:
images = images.to(device)
labels = labels.to(device)
# zero out gradients
optimizer.zero_grad()
# forward
# track when in training
with torch.set_grad_enabled(phase == 'train'):
outs = model(images)
loss = criterion(outs, labels)
_, preds = torch.max(outs.data, 1)
# backward & optimize when in training
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item()
running_correct += (preds == labels).sum().item()
epoch_loss = running_loss / len(dataloader.dataset)
epoch_acc = running_correct / len(dataloader.dataset)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
if phase == 'train' and epoch_acc > best_train_acc:
best_train_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
if phase == 'val':
val_acc_history.append(epoch_acc)
if epoch_acc > best_val_acc:
best_val_acc = epoch_acc
time_elapsed = time.time() - starting_time
time_elapsed = "{:.0f}m {:.0f}s".format(time_elapsed // 60, time_elapsed % 60)
print('Training complete in {}'.format(time_elapsed))
model.load_state_dict(best_model_wts)
return model, best_train_acc, best_val_acc, val_acc_history, time_elapsed
If you need more code, let me know - kinda hard to share as I tried to built it modularly to reuse parts of it.
Thanks!
The issue I could think of, is that the Imagenette Validation set is actually part of the Imagenet Train Set
