def load():
loader = Loader(
mesh2meters_file=args.mesh2meters,
data_dir=args.data_dir,
image_dir=args.image_dir,
)
loader.build_dataset(file="valSeen_data.json", args=args)
loader.build_dataset(file="valUnseen_data.json", args=args)
return loader
batch_size=batch_size,
sampler=train_sampler)
tune_iterator = DataLoader(dataset,
batch_size=batch_size,
sampler=tune_sampler)
return train_iterator, tune_iterator
if __name__ == '__main__':
# load data
loader = Loader(data_dir=args.data_dir,
image_dir=args.image_dir,
target_dir=args.target_dir)
loader.build_dataset(file='train.json',
gaussian_target=args.gaussian_target,
sample_used=args.sample_used)
loader.build_dataset(file='dev.json',
gaussian_target=args.gaussian_target,
sample_used=args.sample_used)
loader.build_dataset(file='test.json',
gaussian_target=args.gaussian_target,
sample_used=args.sample_used)
rnn_args = {
'input_size': len(loader.vocab),
'embed_size': args.embed_size,
'rnn_hidden_size': args.rnn_hidden_size,
'num_rnn_layers': args.num_rnn_layers,
class LingUNetAgent:
def __init__(self, args):
self.args = args
self.device = (torch.device(f"cuda:{args.cuda}")
if torch.cuda.is_available() else torch.device("cpu"))
self.args.device = self.device
self.loss_func = nn.KLDivLoss(reduction="batchmean")
self.loader = None
self.scan_graphs = {}
self.writer = None
self.out_dir = os.path.join(args.log_dir, args.run_name)
if (args.log and args.train) or args.save:
if not os.path.isdir(self.out_dir):
print("Log directory under {}".format(self.out_dir))
os.system("mkdir {}".format(self.out_dir))
self.writer = SummaryWriter(args.summary_dir + args.name)
self.model = None
self.optimizer = None
self.node2pix = json.load(
open(args.image_dir + "allScans_Node2pix.json"))
def testEpochMultiFloor(self, data_iterator, mode, epoch=0):
self.model.eval()
distances = []
with torch.no_grad():
for batch_data in tqdm.tqdm(data_iterator):
(
_,
info_elem,
batch_images,
batch_texts,
batch_seq_lengths,
batch_locations,
_,
all_maps,
all_conversions,
) = batch_data
# get predictions, calculate loss
batch_texts = tile(batch_texts,
dim=0,
n_tile=self.args.max_floors)
batch_seq_lengths = tile(batch_seq_lengths,
dim=0,
n_tile=self.args.max_floors)
preds = self.model(
all_maps.view((
-1,
all_maps.size()[2],
all_maps.size()[3],
all_maps.size()[4],
)).to(device=self.args.device),
batch_texts.to(device=self.args.device),
batch_seq_lengths.to(device=self.args.device),
).view(
all_maps.size(0),
self.args.max_floors,
self.args.ds_height,
self.args.ds_width,
)
# calculate location error and accuracy
if "test" not in mode:
if self.args.distance_metric == "geodesic":
distances.extend(
geo_dist_multifloor(
self.node2pix,
self.scan_graphs[mode],
preds,
all_conversions,
info_elem,
))
elif self.args.distance_metric == "euclidean":
_, _, H, W = batch_images.size()
distances.extend(
euclidean_dist_multifloor(
self.node2pix,
preds,
batch_locations,
all_conversions,
info_elem,
H,
W,
))
if "test" not in mode:
avg_dist = np.mean(distances)
avg_acc = accuracy(distances, 3)
log(self.args, mode, (avg_dist, avg_acc))
else:
format_preds_multifloor_output(
self.node2pix,
self.scan_graphs[mode],
preds,
all_conversions,
info_elem,
self.args.predictions_dir,
mode,
)
return None, None, None
return None, avg_acc, distances
def testEpochSingleFloor(self, data_iterator, mode, epoch=0):
self.model.eval()
distances, total_loss = [], []
with torch.no_grad():
for enum, batch_data in enumerate(tqdm.tqdm(data_iterator)):
(
housemaps,
info_elem,
batch_images,
batch_texts,
batch_seq_lengths,
batch_locations,
batch_conversions,
_,
_,
) = batch_data
# downsample batch locations
_, _, H, W = batch_images.size()
batch_locations = (nn.functional.interpolate(
batch_locations.unsqueeze(1),
(self.args.ds_height, self.args.ds_width),
mode="bilinear",
align_corners=True,
).squeeze(1).float()).to(device=self.args.device)
# get predictions
preds = self.model(
batch_images.to(device=self.args.device),
batch_texts.to(device=self.args.device),
batch_seq_lengths.to(device=self.args.device),
)
if "test" not in mode:
if self.args.distance_metric == "geodesic":
distances.extend(
geo_dist_singlefloor(
self.node2pix,
self.scan_graphs[mode],
preds,
batch_conversions,
info_elem,
))
elif self.args.distance_metric == "euclidean":
distances.extend(
euclidean_dist_singlefloor(preds, batch_locations,
batch_conversions, H,
W))
if self.args.train == True:
loss = self.loss_func(preds, add_epsilon(batch_locations))
total_loss.append(loss.item())
if self.args.visualize and enum == 0:
visualize(
args,
info_elem,
housemaps,
batch_locations,
preds,
distances,
)
if self.args.train:
avg_acc = accuracy(distances, 3)
mean_dist = np.mean(distances)
avg_loss = np.mean(total_loss)
log(self.args, mode, (epoch, avg_loss, mean_dist, avg_acc))
return avg_loss, avg_acc, distances
else:
if "test" not in mode:
avg_acc = accuracy(distances, 3)
mean_dist = np.mean(distances)
log(self.args, mode, (mean_dist, avg_acc))
return None, avg_acc, distances
else:
format_preds_singlefloor_output(
self.node2pix,
self.scan_graphs[mode],
preds,
batch_conversions,
info_elem,
self.args.predictions_dir,
mode,
)
return None, None, None
def run_test(self):
print("Starting Evaluation...")
oldArgs, rnn_args, cnn_args, out_layer_args, state_dict = torch.load(
self.args.eval_ckpt).values()
self.args = load_oldArgs(self.args, oldArgs)
self.load_data()
self.model = LingUNet(rnn_args, cnn_args, self.args)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.model = nn.DataParallel(self.model)
self.model.load_state_dict(state_dict)
self.model = self.model.to(device=self.args.device)
if self.args.test_multi_floor:
_, _, valseen_errors = self.testEpochMultiFloor(
self.valseen_iterator, mode="valSeen")
_, _, valunseen_errors = self.testEpochMultiFloor(
self.val_unseen_iterator, mode="valUnseen")
self.testEpochMultiFloor(self.test_iterator, mode="test")
else:
_, _, valseen_errors = self.testEpochSingleFloor(
self.valseen_iterator, mode="valSeen")
_, _, valunseen_errors = self.testEpochSingleFloor(
self.val_unseen_iterator, mode="valUnseen")
self.testEpochSingleFloor(self.test_iterator, mode="test")
eval_preds({"valSeen": valseen_errors, "valUnseen": valunseen_errors})
def trainEpoch(self):
self.model.train()
total_loss, total_accuracy = [], []
for (
_,
info_elem,
batch_images,
batch_texts,
batch_seq_lengths,
batch_locations,
batch_conversions,
_,
_,
) in tqdm.tqdm(self.train_iterator):
self.optimizer.zero_grad()
# downsample batch locations
_, _, H, W = batch_images.size()
if self.args.data_aug:
crop_size = self.args.ds_height_crop, self.args.ds_width_crop
else:
crop_size = self.args.ds_height, self.args.ds_width
batch_locations = (nn.functional.interpolate(
batch_locations.unsqueeze(1), (crop_size),
mode="bilinear").squeeze(1).float()).to(
device=self.args.device)
# get predictions
preds = self.model(
batch_images.to(device=self.args.device),
batch_texts.to(device=self.args.device),
batch_seq_lengths.to(device=self.args.device),
)
# calculate loss
loss = self.loss_func(preds, add_epsilon(batch_locations))
total_loss.append(loss.item())
loss.backward()
# calculate location error and accuracy
if self.args.distance_metric == "geodesic":
total_accuracy.append(
accuracy(
geo_dist_singlefloor(
self.node2pix,
self.scan_graphs["train"],
preds,
batch_conversions,
info_elem,
),
3,
))
elif self.args.distance_metric == "euclidean":
total_accuracy.append(
accuracy(
euclidean_dist_singlefloor(preds, batch_locations,
batch_conversions, H, W),
3,
))
nn.utils.clip_grad_value_(self.model.parameters(),
clip_value=self.args.grad_clip)
self.optimizer.step()
train_loss = np.mean(total_loss)
train_acc = np.mean(np.asarray(total_accuracy))
return train_acc, train_loss
def run_train(self):
assert self.args.num_lingunet_layers is not None
rnn_args = {"input_size": len(self.loader.vocab)}
cnn_args = {
"kernel_size": 5,
"padding": 2,
"deconv_dropout": 0,
"conv_dropout": 0,
}
self.model = LingUNet(rnn_args, cnn_args, args)
num_params = sum(
[p.numel() for p in self.model.parameters() if p.requires_grad])
print("Number of parameters:", num_params)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.model = nn.DataParallel(self.model)
self.model = self.model.to(device=self.args.device)
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.args.lr)
print("Starting Training...")
best_tune_acc = float("-inf")
best_unseen_acc = float("-inf")
best_model = None
patience = 0
save_path = ""
for epoch in range(self.args.num_epoch):
train_acc, train_loss = self.trainEpoch()
valseen_loss, valseen_acc, _ = self.testEpochSingleFloor(
self.valseen_iterator, "valSeen", epoch)
val_unseen_loss, val_unseen_acc, _ = self.testEpochSingleFloor(
self.val_unseen_iterator, "valUnseen", epoch)
if self.args.log:
self.writer.add_scalar("Accuracy/train", train_acc, epoch)
self.writer.add_scalar("Loss/train", train_loss, epoch)
self.writer.add_scalar("Accuracy/val_seen", valseen_acc, epoch)
self.writer.add_scalar("Loss/val_seen", valseen_loss, epoch)
self.writer.add_scalar("Accuracy/val_unseen", val_unseen_acc,
epoch)
self.writer.add_scalar("Loss/val_unseen", val_unseen_loss,
epoch)
if valseen_acc > best_tune_acc or val_unseen_acc > best_unseen_acc:
best_model = copy.deepcopy(self.model)
if self.args.save:
save_path = os.path.join(
self.out_dir,
"{}_acc{:.2f}_unseenAcc{:.2f}_epoch{}.pt".format(
self.args.name, valseen_acc, val_unseen_acc,
epoch),
)
state = convert_model_to_state(best_model, args, rnn_args,
cnn_args)
torch.save(state, save_path)
if valseen_acc > best_tune_acc:
best_tune_acc = valseen_acc
print("[Tune]: Best valSeen accuracy:", best_tune_acc)
if val_unseen_acc > best_unseen_acc:
best_unseen_acc = val_unseen_acc
print("[Tune]: Best valUNseen accuracy:", best_unseen_acc)
patience = 0
else:
patience += 1
# if patience reachs threshold end the training
if patience >= self.args.early_stopping:
break
print("Patience:", patience)
print(f"Best model saved at: {save_path}")
def load_data(self):
self.loader = Loader(
mesh2meters_file=self.args.mesh2meters,
data_dir=self.args.data_dir,
image_dir=self.args.image_dir,
)
if self.args.train and self.args.increase_train:
self.loader.build_dataset_extra_train(file="train_data.json",
args=self.args)
else:
self.loader.build_dataset(file="train_data.json", args=self.args)
self.loader.build_dataset(file="valSeen_data.json", args=self.args)
self.loader.build_dataset(file="valUnseen_data.json", args=self.args)
self.train_iterator = DataLoader(
self.loader.datasets["train"],
batch_size=self.args.batch_size,
shuffle=True,
)
self.valseen_iterator = DataLoader(
self.loader.datasets["valSeen"],
batch_size=self.args.batch_size,
shuffle=False,
)
self.val_unseen_iterator = DataLoader(
self.loader.datasets["valUnseen"],
batch_size=self.args.batch_size,
shuffle=False,
)
self.scan_graphs["train"] = self.loader.datasets["train"].scan_graphs
self.scan_graphs["valSeen"] = self.loader.datasets[
"valSeen"].scan_graphs
self.scan_graphs["valUnseen"] = self.loader.datasets[
"valUnseen"].scan_graphs
if self.args.evaluate:
self.loader.build_dataset(file="test_data.json", args=self.args)
self.test_iterator = DataLoader(
self.loader.datasets["test"],
batch_size=self.args.batch_size,
shuffle=False,
)
self.scan_graphs["test"] = self.loader.datasets["test"].scan_graphs
def run(self):
if self.args.train:
self.load_data()
self.run_train()
elif self.args.evaluate:
self.run_test()
