def _reconstruction_errors(images):
latent = encoder(images, training=False) # E(x)
reconstructed_images = generator(latent, training=False) # G(E(x))
features = discriminator_features([images, latent], training=False) # f_D(x, E(x))
reconstructed_features = discriminator_features([reconstructed_images, latent], training=False) # f_D(G(E(x)), E(x))
pixel_distance = l1(images, reconstructed_images) # L_R
features_distance = l1(features, reconstructed_features) # L_f_D
return pixel_distance, features_distance
def validate(self, epoch, output_save):
self.model.eval()
batch_loss = 0.0
batch_iou = 0.0
vis_save = os.path.join(output_save, "epoch%02d" % (epoch+1))
n_batches = len(self.dataloader_val)
with torch.no_grad():
for idx, sample in enumerate(self.dataloader_val):
input = sample['occ_grid'].to(self.device)
target_df = sample['df_gt'].to(self.device)
names = sample['name']
# ===================forward=====================
output_df = self.model(input)
loss = losses.l1(output_df, target_df, use_log_transform=False)
iou = metric.iou_df(output_df, target_df, trunc_dist=1.0)
# ===================log========================
batch_loss += loss.item()
batch_iou += iou
# save the predictions at the end of the epoch
# if epoch > args.save_epoch and (idx + 1) == n_batches-1:
# pred_dfs = output_df[:args.n_vis + 1]
# target_dfs = target_df[:args.n_vis + 1]
# names = names[:args.n_vis + 1]
# utils.save_predictions(vis_save, args.model_name, args.gt_type, names, pred_dfs=pred_dfs, target_dfs=target_dfs,
# pred_occs=None, target_occs=None)
val_loss = batch_loss / (idx + 1)
mean_iou = batch_iou / (idx + 1)
return val_loss, mean_iou
def train(self, epoch):
self.model.train()
batch_loss = 0.0
for idx, sample in enumerate(self.dataloader_train):
input = sample['occ_grid'].to(self.device)
target = sample['df_gt'].to(self.device)
# zero the parameter gradients
self.optimizer.zero_grad()
# ===================forward=====================
output = self.model(input)
loss = losses.l1(output, target, use_log_transform=args.use_logweight)
# ===================backward + optimize====================
loss.backward()
self.optimizer.step()
# ===================log========================
batch_loss += loss.item()
# if (idx + 1) % 10 == 0:
# print('Training : [iter %d / epoch %d] loss: %.3f' % (idx + 1, epoch + 1, loss.item()))
train_loss = batch_loss / (idx + 1)
return train_loss
def validate(self, epoch, output_save):
self.model.eval()
batch_loss_bce = 0.0
batch_loss_l1 = 0.0
batch_iou = 0.0
vis_save = os.path.join(output_save, "epoch%02d" % (epoch + 1))
n_batches = len(self.dataloader_val)
with torch.no_grad():
for idx, sample in enumerate(self.dataloader_val):
input = sample['occ_grid'].to(self.device)
target_occ = sample['occ_gt'].to(self.device)
target_df = sample['occ_df_gt'].to(self.device)
names = sample['name']
# ===================forward=====================
output_occ = self.model(input)
loss_bce = losses.bce(output_occ, target_occ)
# Convert occ to df to calculate l1 loss
output_df = utils.occs_to_dfs(output_occ,
trunc=args.truncation,
pred=True)
loss_l1 = losses.l1(output_df, target_df)
iou = metric.iou_occ(output_occ, target_occ)
# ===================log========================
batch_loss_bce += loss_bce.item()
batch_loss_l1 += loss_l1.item()
batch_iou += iou
# save the predictions at the end of the epoch
# if epoch > args.save_epoch and (idx + 1) == n_batches-1:
# pred_occs = output_occ[:args.n_vis+1]
# target_occs = target_occ[:args.n_vis+1]
# names = names[:args.n_vis+1]
# utils.save_predictions(vis_save, args.model_name, args.gt_type, names, pred_dfs=None, target_dfs=None,
# pred_occs=pred_occs, target_occs=target_occs)
val_loss_bce = batch_loss_bce / (idx + 1)
val_loss_l1 = batch_loss_l1 / (idx + 1)
mean_iou = batch_iou / (idx + 1)
return val_loss_bce, val_loss_l1, mean_iou
def test(test_list):
dataset_test = dataloader.DatasetLoad(data_list=test_list,
truncation=args.truncation)
if args.model_name == 'Net3D':
model = Net3D(1, 1).to(device)
elif args.model_name == 'UNet3D':
model = UNet3D(1, 1).to(device)
dataloader_test = torchdata.DataLoader(dataset_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=2,
drop_last=False)
# load our saved model and use it to predict the class for test images
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'])
model = model.to(device)
model.eval()
vis_save = "%sfinal_results/vis/%s/%s" % (params["network_output"],
args.model_name, args.gt_type)
pathlib.Path(vis_save).mkdir(parents=True, exist_ok=True)
batch_l1 = 0.0
batch_iou = 0.0
n_batches = len(dataloader_test)
with torch.no_grad():
for idx, sample in enumerate(dataloader_test):
input = sample['occ_grid'].to(device)
names = sample['name']
if args.gt_type == 'occ':
target_occ = sample['occ_gt'].to(device)
target_df = sample['occ_df_gt'].to(device)
# ===================forward=====================
output_occ = model(input)
# Convert occ to df to calculate l1 loss
output_df = utils.occs_to_dfs(output_occ,
trunc=args.truncation,
pred=True)
l1 = losses.l1(output_df, target_df)
iou = metric.iou_occ(output_occ, target_occ)
# save the predictions
if (idx + 1) > n_batches - 2:
pred_occs = output_occ[:args.n_vis + 1]
target_occs = target_occ[:args.n_vis + 1]
names = names[:args.n_vis + 1]
utils.save_predictions(vis_save,
args.model_name,
args.gt_type,
names,
pred_dfs=None,
target_dfs=None,
pred_occs=pred_occs,
target_occs=target_occs)
else:
target_df = sample['df_gt'].to(device)
output_df = model(input)
l1 = losses.l1(output_df, target_df)
iou = metric.iou_df(output_df, target_df, trunc_dist=1.0)
# save the predictions
if (idx + 1) > n_batches - 2:
pred_dfs = output_df[:args.n_vis + 1]
target_dfs = target_df[:args.n_vis + 1]
names = names[:args.n_vis + 1]
utils.save_predictions(vis_save,
args.model_name,
args.gt_type,
names,
pred_dfs=pred_dfs,
target_dfs=target_dfs,
pred_occs=None,
target_occs=None)
batch_l1 += l1.item()
batch_iou += iou
l1_error = batch_l1 / (idx + 1)
mean_iou = batch_iou / (idx + 1)
print("Mean IOU: ", mean_iou)
print("L1 Error: ", l1_error)