## cpu dtype
dtype = torch.FloatTensor
save_model_path = "model_state_dict.pkl"
csv_path = '../../data/train_v2.csv'
img_path = '../../data/train-jpg'
img_ext = '.jpg'
transform_list = [transforms.Scale(224)]
training_dataset = AmazonDataset(csv_path,
img_path,
img_ext,
dtype,
transform_list=transform_list,
three_band=True)
## loader
train_loader = DataLoader(
training_dataset,
batch_size=128,
shuffle=True,
num_workers=4, # 0 for CUDA
)
model = torchvision.models.vgg11(pretrained=True)
model.type(dtype)
loss_fn = nn.MultiLabelSoftMarginLoss().type(dtype)
acc, loss = validate_epoch(model, train_loader, loss_fn, dtype)
print(acc, loss)
acc_history = []
loss_history = []
## don't load model params from file - instead retrain the model
if not from_pickle:
for epoch in range(num_epochs):
print("Begin epoch {}/{}".format(epoch + 1, num_epochs))
epoch_loss = train_epoch(train_loader,
model,
loss_fn,
optimizer,
dtype,
print_every=10)
scheduler.step(epoch_loss[-1], epoch)
## f2 score for validation dataset
acc = validate_epoch(model, val_loader, dtype)
acc_history.append(acc)
loss_history += epoch_loss
print("END epoch {}/{}: F2 score = {:.02f}".format(
epoch + 1, num_epochs, acc))
## serialize model data and save as .pkl file
torch.save(model.state_dict(), save_model_path)
print("model saved as {}".format(os.path.abspath(save_model_path)))
## save loss and accuracy as .mat file
savemat(save_mat_path, {
"acc": acc_history,
"loss": loss_history,
"num_epochs": num_epochs
})
## load model params from file
else:
factor=adaptive_lr_factor)
for epoch in range(1, num_epochs + 1):
print("Begin epoch {}/{}".format(epoch, num_epochs))
epoch_losses, epoch_f2 = train_epoch(
train_loader,
model,
loss_fn,
optimizer,
dtype,
sigmoid_threshold=sigmoid_threshold,
print_every=20)
scheduler.step(np.mean(epoch_losses), epoch)
## f2 score for validation dataset
f2_acc = validate_epoch(model,
val_loader,
dtype,
sigmoid_threshold=sigmoid_threshold)
## store results
train_acc_history += epoch_f2
val_acc_history.append(f2_acc)
loss_history += epoch_losses
## overwrite the model .pkl file every epoch
torch.save(model.state_dict(), save_model_path)
save_accuracy_and_loss_mat(save_mat_path,
train_acc_history,
val_acc_history,
loss_history,
epoch,
lr=optimizer.param_groups[-1]['lr'])
## checkpoints
if save_every and not epoch % save_every:
for param in model.parameters():
param.requires_grad = False
for param in model.fc.parameters():
param.requires_grad = True
for param in model.conv1.parameters():
param.requires_grad = True
loss_fn = nn.MultiLabelSoftMarginLoss().type(dtype)
optimizer = optim.Adam(chain(model.fc.parameters(),
model.conv1.parameters()),
lr=1e-3)
## don't load model params from file - instead retrain the model
if not from_pickle:
train(train_loader, model, loss_fn, optimizer, dtype, print_every=1)
## serialize model data and save as .pkl file
torch.save(model.state_dict(), save_model_path)
print("model saved as {}".format(os.path.abspath(save_model_path)))
## load model params from file
else:
state_dict = torch.load(save_model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
print("model loaded from {}".format(os.path.abspath(save_model_path)))
train_acc_loader = DataLoader(training_dataset,
batch_size=200,
shuffle=True,
num_workers=6)
acc = validate_epoch(model, train_acc_loader, dtype)
print(acc)