def train_model(model,
pretrained_features,
target_column,
labels_file,
train_images_folder,
valid_images_folder=None,
batch_size=32,
num_workers=4,
num_epochs=10,
use_gpu=None,
flatten_pretrained_out=False):
train_dset_loader = make_dsets(train_images_folder,
labels_file,
target_column,
batch_size=batch_size,
num_workers=num_workers,
is_train=True)
valid_dset_loader = None
if valid_images_folder:
valid_dset_loader = make_dsets(valid_images_folder,
labels_file,
target_column,
batch_size=batch_size,
num_workers=num_workers,
is_train=False)
# Sleeve Length Model
model = train_attribute_model(
model,
pretrained_features,
train_dset_loader=train_dset_loader,
valid_dset_loader=valid_dset_loader,
num_epochs=num_epochs,
use_gpu=use_gpu,
flatten_pretrained_out=flatten_pretrained_out)
return model
def evaluate_model(model,
pretrained_model,
target_column,
labels_file,
images_folder,
batch_size=32,
num_workers=4,
use_gpu=None,
flatten_pretrained_out=False):
if use_gpu is None:
use_gpu = torch.cuda.is_available()
dset_loader = make_dsets(images_folder,
labels_file,
target_column,
is_train=False,
shuffle=False,
batch_size=batch_size,
num_workers=num_workers)
running_loss, running_corrects = (0, 0)
y_true = []
y_pred = []
y_pred_proba = []
input_files = []
model.eval()
for batch_idx, data in enumerate(dset_loader):
# Get the inputs
batch_files, inputs, labels = data
y_true = np.concatenate([y_true, labels.numpy()])
input_files = np.concatenate([input_files, batch_files])
# Wrap them in Variable
if use_gpu:
inputs, labels = inputs.cuda(), labels.cuda()
inputs = Variable(inputs)
labels = Variable(labels)
# Get output from pre-trained model and re-shape to Flatten
out_features = pretrained_model(inputs)
if flatten_pretrained_out:
out_features = out_features.view(out_features.size(0), -1)
# Forward
outputs = model(out_features)
loss = F.nll_loss(outputs, labels)
preds_proba, preds = outputs.data.max(1)
y_pred = np.concatenate([y_pred, preds.cpu().numpy().flatten()])
y_pred_proba = np.concatenate(
[y_pred_proba, preds_proba.cpu().numpy().flatten()])
# Statistics
running_loss += loss.data[0]
running_corrects += torch.sum(preds == labels.data)
return {
"loss": running_loss,
"accuracy": running_corrects / len(y_true),
"input_files": input_files,
"y_true": y_true,
"y_pred": y_pred,
"y_pred_proba": y_pred_proba
}