def main(args):
data_config = load_config_from_json(args.data_config_path)
model_config = load_config_from_json(
os.path.join(args.saved_model_path, "config.jsonl"))
# initialize model
model = SFNet(model_config["sfnet"])
model = model.to(device)
if not os.path.exists(args.saved_model_path):
raise FileNotFoundError(args.saved_model_path)
checkpoint = os.path.join(args.saved_model_path, args.checkpoint)
model.load_state_dict(torch.load(checkpoint, map_location="cpu"))
print("Model loaded from %s" % (args.saved_model_path))
# tracker to keep true labels and predicted probabilitites
target_tracker = []
pred_tracker = []
print("Preparing test data ...")
dataset = ModCloth(data_config, split="test")
data_loader = DataLoader(
dataset=dataset,
batch_size=model_config["trainer"]["batch_size"],
shuffle=False,
)
print("Evaluating model on test data ...")
model.eval()
with torch.no_grad():
for iteration, batch in enumerate(data_loader):
for k, v in batch.items():
if torch.is_tensor(v):
batch[k] = to_var(v)
# Forward pass
_, pred_probs = model(batch)
target_tracker.append(batch["fit"].cpu().numpy())
pred_tracker.append(pred_probs.cpu().data.numpy())
target_tracker = np.stack(target_tracker[:-1]).reshape(-1)
pred_tracker = np.stack(pred_tracker[:-1], axis=0).reshape(
-1, model_config["sfnet"]["num_targets"])
precision, recall, f1_score, accuracy, auc = compute_metrics(
target_tracker, pred_tracker)
print("-" * 50)
print(
"Metrics:\n Precision = {:.3f}\n Recall = {:.3f}\n F1-score = {:.3f}\n Accuracy = {:.3f}\n AUC = {:.3f}\n "
.format(precision, recall, f1_score, accuracy, auc))
print("-" * 50)
print("Instantiate dataloader")
test_dataset = PF_Pascal(args.test_csv_path, args.test_image_path,
args.feature_h, args.feature_w, args.eval_type)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=1,
shuffle=False,
num_workers=args.num_workers)
# Instantiate model
print("Instantiate model")
net = SFNet(args.feature_h,
args.feature_w,
beta=args.beta,
kernel_sigma=args.kernel_sigma)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# Load weights
print("Load pre-trained weights")
best_weights = torch.load("./weights/best_checkpoint.pt")
adap3_dict = best_weights['state_dict1']
adap4_dict = best_weights['state_dict2']
net.adap_layer_feat3.load_state_dict(adap3_dict, strict=False)
net.adap_layer_feat4.load_state_dict(adap4_dict, strict=False)
# PCK metric from 'https://github.com/ignacio-rocco/weakalign/blob/master/util/eval_util.py'
def correct_keypoints(source_points, warped_points, L_pck, alpha=0.1):
# compute correct keypoints
p_src = source_points[0, :]
p_wrp = warped_points[0, :]
def main(args):
ts = time.strftime("%Y-%b-%d-%H-%M-%S", time.gmtime())
data_config = load_config_from_json(args.data_config_path)
model_config = load_config_from_json(args.model_config_path)
splits = ["train", "valid"]
datasets = OrderedDict()
for split in splits:
datasets[split] = ModCloth(data_config, split=split)
# initialize model
model = SFNet(model_config["sfnet"])
model = model.to(device)
print("-" * 50)
print(model)
print("-" * 50)
print("Number of model parameters: {}".format(
sum(p.numel() for p in model.parameters())))
print("-" * 50)
save_model_path = os.path.join(
model_config["logging"]["save_model_path"],
model_config["logging"]["run_name"] + ts,
)
os.makedirs(save_model_path)
if model_config["logging"]["tensorboard"]:
writer = SummaryWriter(os.path.join(save_model_path, "logs"))
writer.add_text("model", str(model))
writer.add_text("args", str(args))
loss_criterion = torch.nn.CrossEntropyLoss(reduction="mean")
optimizer = torch.optim.Adam(
model.parameters(),
lr=model_config["trainer"]["optimizer"]["lr"],
weight_decay=model_config["trainer"]["optimizer"]["weight_decay"],
)
step = 0
tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.Tensor
for epoch in range(model_config["trainer"]["num_epochs"]):
for split in splits:
data_loader = DataLoader(
dataset=datasets[split],
batch_size=model_config["trainer"]["batch_size"],
shuffle=split == "train",
)
loss_tracker = defaultdict(tensor)
# Enable/Disable Dropout
if split == "train":
model.train()
else:
model.eval()
target_tracker = []
pred_tracker = []
for iteration, batch in enumerate(data_loader):
for k, v in batch.items():
if torch.is_tensor(v):
batch[k] = to_var(v)
# Forward pass
logits, pred_probs = model(batch)
# loss calculation
loss = loss_criterion(logits, batch["fit"])
# backward + optimization
if split == "train":
optimizer.zero_grad()
loss.backward()
optimizer.step()
step += 1
# bookkeepeing
loss_tracker["Total Loss"] = torch.cat(
(loss_tracker["Total Loss"], loss.view(1)))
if model_config["logging"]["tensorboard"]:
writer.add_scalar(
"%s/Total Loss" % split.upper(),
loss.item(),
epoch * len(data_loader) + iteration,
)
if iteration % model_config["logging"][
"print_every"] == 0 or iteration + 1 == len(
data_loader):
print("{} Batch Stats {}/{}, Loss={:.2f}".format(
split.upper(), iteration,
len(data_loader) - 1, loss.item()))
if split == "valid":
target_tracker.append(batch["fit"].cpu().numpy())
pred_tracker.append(pred_probs.cpu().data.numpy())
print("%s Epoch %02d/%i, Mean Total Loss %9.4f" % (
split.upper(),
epoch + 1,
model_config["trainer"]["num_epochs"],
torch.mean(loss_tracker["Total Loss"]),
))
if model_config["logging"]["tensorboard"]:
writer.add_scalar(
"%s-Epoch/Total Loss" % split.upper(),
torch.mean(loss_tracker["Total Loss"]),
epoch,
)
# Save checkpoint
if split == "train":
checkpoint_path = os.path.join(save_model_path,
"E%i.pytorch" % (epoch + 1))
torch.save(model.state_dict(), checkpoint_path)
print("Model saved at %s" % checkpoint_path)
if split == "valid" and model_config["logging"]["tensorboard"]:
# not considering the last (incomplete) batch for metrics
target_tracker = np.stack(target_tracker[:-1]).reshape(-1)
pred_tracker = np.stack(pred_tracker[:-1], axis=0).reshape(
-1, model_config["sfnet"]["num_targets"])
precision, recall, f1_score, accuracy, auc = compute_metrics(
target_tracker, pred_tracker)
writer.add_scalar("%s-Epoch/Precision" % split.upper(), precision,
epoch)
writer.add_scalar("%s-Epoch/Recall" % split.upper(), recall, epoch)
writer.add_scalar("%s-Epoch/F1-Score" % split.upper(), f1_score,
epoch)
writer.add_scalar("%s-Epoch/Accuracy" % split.upper(), accuracy,
epoch)
writer.add_scalar("%s-Epoch/AUC" % split.upper(), auc, epoch)
# Save Model Config File
with jsonlines.open(os.path.join(save_model_path, "config.jsonl"),
"w") as fout:
fout.write(model_config)