def cal_l2_losses(pred_traj_gt, pred_traj_gt_rel, pred_traj_fake,
pred_traj_fake_rel, loss_mask):
g_l2_loss_abs = l2_loss(pred_traj_fake,
pred_traj_gt,
loss_mask,
mode='sum')
g_l2_loss_rel = l2_loss(pred_traj_fake_rel,
pred_traj_gt_rel,
loss_mask,
mode='sum')
return g_l2_loss_abs, g_l2_loss_rel
def generator_step(args, batch, generator, optimizer_g):
# batch = [tensor.cuda() for tensor in batch]
(obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, non_linear_ped,
loss_mask, seq_start_end) = batch
losses = {}
loss = torch.zeros(1).to(pred_traj_gt)
print("loss", loss)
loss = Variable(loss)
g_l2_loss_rel = []
loss_mask = loss_mask[:, args.obs_len:]
# for _ in range(args.best_k):
obs_traj = obs_traj.float()
obs_traj_rel = obs_traj_rel.float()
generator_out = generator(obs_traj, obs_traj_rel, seq_start_end)
pred_traj_fake_rel = generator_out
# pred_traj_fake = relative_to_abs(pred_traj_fake_rel, obs_traj[-1])
if args.l2_loss_weight > 0:
g_l2_loss_rel.append(args.l2_loss_weight * l2_loss(
pred_traj_fake_rel, pred_traj_gt_rel, loss_mask, mode='raw'))
g_l2_loss_sum_rel = torch.zeros(1).to(pred_traj_gt)
if args.l2_loss_weight > 0:
g_l2_loss_rel = torch.stack(g_l2_loss_rel, dim=1)
for start, end in seq_start_end.data:
_g_l2_loss_rel = g_l2_loss_rel[start:end]
_g_l2_loss_rel = torch.sum(_g_l2_loss_rel, dim=0)
_g_l2_loss_rel = torch.min(_g_l2_loss_rel) / torch.sum(
loss_mask[start:end])
g_l2_loss_sum_rel += _g_l2_loss_rel
losses['G_l2_loss_rel'] = g_l2_loss_sum_rel.item()
loss += g_l2_loss_sum_rel
# traj_fake = torch.cat([obs_traj, pred_traj_fake], dim=0)
traj_fake_rel = torch.cat([obs_traj_rel, pred_traj_fake_rel], dim=0)
# scores_fake = discriminator(traj_fake, traj_fake_rel, seq_start_end)
# discriminator_loss = g_loss_fn(scores_fake)
# loss += discriminator_loss
# losses['G_discriminator_loss'] = discriminator_loss.item()
losses['G_total_loss'] = loss.item()
# loss = Variable(loss, requires_grad= True)
# optimizer_g.zero_grad()
loss.backward()
# if args.clipping_threshold_g > 0:
# nn.utils.clip_grad_norm_(
# generator.parameters(), args.clipping_threshold_g
# )
optimizer_g.step()
optimizer_g.zero_grad()
return losses
def testL2Loss(self):
with self.test_session():
shape = [5, 5, 5]
num_elem = 5 * 5 * 5
weights = tf.constant(1.0, shape=shape)
wd = 0.01
loss = losses.l2_loss(weights, wd)
self.assertEquals(loss.op.name, 'L2Loss/value')
self.assertAlmostEqual(loss.eval(), num_elem * wd / 2, 5)
def generator_step(args, batch, generator, discriminator, g_loss_fn,
optimizer_g, discriminator_wight):
if args.use_gpu == 1:
batch = [tensor.cuda() for tensor in batch]
else:
batch = [tensor for tensor in batch]
# batch = [tensor.cuda() for tensor in batch]
(obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, non_linear_ped,
loss_mask, seq_start_end) = batch
losses = {}
loss = torch.zeros(1).to(pred_traj_gt)
g_l2_loss_rel = []
loss_mask = loss_mask[:, args.obs_len:]
for _ in range(args.best_k):
generator_out = generator(obs_traj, obs_traj_rel, seq_start_end)
pred_traj_fake_rel = generator_out
pred_traj_fake = relative_to_abs(pred_traj_fake_rel, obs_traj[-1])
if args.l2_loss_weight > 0:
g_l2_loss_rel.append(args.l2_loss_weight * l2_loss(
pred_traj_fake_rel, pred_traj_gt_rel, loss_mask, mode='raw'))
g_l2_loss_sum_rel = torch.zeros(1).to(pred_traj_gt)
if args.l2_loss_weight > 0:
g_l2_loss_rel = torch.stack(g_l2_loss_rel, dim=1)
for start, end in seq_start_end.data:
_g_l2_loss_rel = g_l2_loss_rel[start:end]
_g_l2_loss_rel = torch.sum(_g_l2_loss_rel, dim=0)
_g_l2_loss_rel = torch.min(_g_l2_loss_rel) / torch.sum(
loss_mask[start:end])
g_l2_loss_sum_rel += _g_l2_loss_rel
losses['G_l2_loss_rel'] = g_l2_loss_sum_rel.item()
loss += g_l2_loss_sum_rel
traj_fake = torch.cat([obs_traj, pred_traj_fake], dim=0)
traj_fake_rel = torch.cat([obs_traj_rel, pred_traj_fake_rel], dim=0)
scores_fake = discriminator(traj_fake, traj_fake_rel, seq_start_end)
discriminator_loss = g_loss_fn(scores_fake)
loss += (discriminator_wight * discriminator_loss)
losses['G_discriminator_loss'] = discriminator_loss.item()
losses['G_total_loss'] = loss.item()
optimizer_g.zero_grad()
loss.backward()
if args.clipping_threshold_g > 0:
nn.utils.clip_grad_norm_(generator.parameters(),
args.clipping_threshold_g)
optimizer_g.step()
return losses
fake_img, fake_mask = generator(real_img, desired_au, reuse=False)
fake_img_masked = fake_mask * real_img + (1 - fake_mask) * fake_img
# G(G(Ic1, c2)*M, c1) * M
cyc_img, cyc_mask = generator(fake_img_masked, real_au, reuse=True)
cyc_img_masked = cyc_mask * fake_img_masked + (1 - cyc_mask) * cyc_img
# D(real_I)
pred_real_img, pred_real_au = discriminator(real_img, reuse=False)
# D(fake_I)
pred_fake_img_masked, pred_fake_au = discriminator(fake_img_masked, reuse=True)
# ============= losses =============
# G losses
loss_g_fake_img_masked = -tf.reduce_mean(pred_fake_img_masked) * lambda_D_img
loss_g_fake_au = l2_loss(desired_au, pred_fake_au) * lambda_D_au
loss_g_cyc = l1_loss(real_img, cyc_img_masked) * lambda_cyc
loss_g_mask_fake = tf.reduce_mean(fake_mask) * lambda_mask + smooth_loss(fake_mask) * lambda_mask_smooth
loss_g_mask_cyc = tf.reduce_mean(cyc_mask) * lambda_mask + smooth_loss(cyc_mask) * lambda_mask_smooth
loss_g = loss_g_fake_img_masked + loss_g_fake_au + \
loss_g_cyc + \
loss_g_mask_fake + loss_g_mask_cyc
# D losses
loss_d_img = -tf.reduce_mean(pred_real_img) * lambda_D_img + tf.reduce_mean(pred_fake_img_masked) * lambda_D_img
loss_d_au = l2_loss(real_au, pred_real_au) * lambda_D_au
alpha = tf.random_uniform([BATCH_SIZE, 1, 1, 1], minval=0., maxval=1.)
differences = fake_img_masked - real_img
def fc(inputs,
num_units_out,
activation=tf.nn.relu,
stddev=0.01,
bias=None,
weight_decay=0,
batch_norm_params=None,
is_training=True,
trainable=True,
restore=True,
scope=None):
"""Adds a fully connected layer followed by an optional batch_norm layer.
FC creates a variable called 'weights', representing the fully connected
weight matrix, that is multiplied by the input. If `batch_norm` is None, a
second variable called 'biases' is added to the result of the initial
vector-matrix multiplication.
Args:
inputs: a [B x N] tensor where B is the batch size and N is the number of
input units in the layer.
num_units_out: the number of output units in the layer.
activation: activation function.
stddev: the standard deviation for the weights.
bias: the initial value of the biases.
weight_decay: the weight decay.
batch_norm_params: parameters for the batch_norm. If is None don't use it.
is_training: whether or not the model is in training mode.
trainable: whether or not the variables should be trainable or not.
restore: whether or not the variables should be marked for restore.
scope: Optional scope for variable_op_scope.
Returns:
the tensor variable representing the result of the series of operations.
"""
with tf.variable_op_scope([inputs], scope, 'FC'):
num_units_in = inputs.get_shape()[1]
weights_shape = [num_units_in, num_units_out]
#weights_initializer = tf.truncated_normal_initializer(stddev=stddev)
weights_initializer = tf.random_normal_initializer(stddev=stddev)
l2_regularizer = lambda t: losses.l2_loss(t, weight_decay)
weights = variables.variable('weights',
shape=weights_shape,
initializer=weights_initializer,
regularizer=l2_regularizer,
trainable=trainable,
restore=restore)
if batch_norm_params is not None:
outputs = tf.matmul(inputs, weights)
with scopes.arg_scope([batch_norm], is_training=is_training,
trainable=trainable, restore=restore):
outputs = batch_norm(outputs, **batch_norm_params)
elif bias is None:
outputs = tf.matmul(inputs, weights)
else:
bias_shape = [num_units_out,]
bias_initializer = tf.constant_initializer(bias)
biases = variables.variable('biases',
shape=bias_shape,
initializer=bias_initializer,
trainable=trainable,
restore=restore)
outputs = tf.nn.xw_plus_b(inputs, weights, biases)
if activation:
outputs = activation(outputs)
return outputs
def conv2d(inputs,
num_filters_out,
kernel_size,
stride=1,
padding='SAME',
activation=tf.nn.relu,
stddev=0.01,
bias=None,
weight_decay=0,
batch_norm_params=None,
is_training=True,
trainable=True,
restore=True,
scope=None):
"""Adds a 2D convolution followed by an optional batch_norm layer.
conv2d creates a variable called 'weights', representing the convolutional
kernel, that is convolved with the input. If `batch_norm_params` is None, a
second variable called 'biases' is added to the result of the convolution
operation.
Args:
inputs: a tensor of size [batch_size, height, width, channels].
num_filters_out: the number of output filters.
kernel_size: a 2-D list comprising of the height and width of the filters.
stride: the stride in height and width of the convolution.
padding: one of 'VALID' or 'SAME'.
activation: activation function.
stddev: standard deviation of the truncated guassian weight distribution.
bias: the initial value of the biases.
weight_decay: the weight decay.
batch_norm_params: parameters for the batch_norm. If is None don't use it.
is_training: whether or not the model is in training mode.
trainable: whether or not the variables should be trainable or not.
restore: whether or not the variables should be marked for restore.
scope: Optional scope for variable_op_scope.
Returns:
a tensor representing the output of the operation.
Raises:
ValueError: if 'kernel_size' is not a 2-D list.
"""
if len(kernel_size) != 2:
raise ValueError('kernel_size must be a 2-D list.')
with tf.variable_op_scope([inputs], scope, 'Conv'):
num_filters_in = inputs.get_shape()[-1]
#num_filters_in = inputs.get_shape()[-3]
weights_shape = [kernel_size[0], kernel_size[1],
num_filters_in, num_filters_out]
#weights_initializer = tf.truncated_normal_initializer(stddev=stddev)
stddev = math.sqrt(2./(kernel_size[0]*kernel_size[1]*num_filters_out))
weights_initializer = tf.random_normal_initializer(stddev=stddev)
l2_regularizer = lambda t: losses.l2_loss(t, weight_decay)
weights = variables.variable('weights',
shape=weights_shape,
initializer=weights_initializer,
regularizer=l2_regularizer,
trainable=trainable,
restore=restore)
conv = tf.nn.conv2d(inputs, weights, [1, stride, stride, 1],
padding=padding)
#conv = tf.nn.conv2d(inputs, weights, [1, 1, stride, stride],
# padding=padding, data_format='NCHW')
if batch_norm_params is not None:
with scopes.arg_scope([batch_norm], is_training=is_training,
trainable=trainable, restore=restore):
outputs = batch_norm(conv, **batch_norm_params)
elif bias is None:
outputs = conv
else:
bias_shape = [num_filters_out,]
bias_initializer = tf.constant_initializer(bias)
biases = variables.variable('biases',
shape=bias_shape,
initializer=bias_initializer,
trainable=trainable,
restore=restore)
outputs = tf.nn.bias_add(conv, biases)
#outputs = tf.nn.bias_add(conv, biases, data_format='NCHW')
#outputs = tf.reshape(outputs, conv.get_shape())
if activation:
outputs = activation(outputs)
return outputs
def fc(inputs,
num_units_out,
activation=tf.nn.relu,
stddev=0.01,
bias=None,
weight_decay=0,
batch_norm_params=None,
is_training=True,
trainable=True,
restore=True,
scope=None):
"""Adds a fully connected layer followed by an optional batch_norm layer.
FC creates a variable called 'weights', representing the fully connected
weight matrix, that is multiplied by the input. If `batch_norm` is None, a
second variable called 'biases' is added to the result of the initial
vector-matrix multiplication.
Args:
inputs: a [B x N] tensor where B is the batch size and N is the number of
input units in the layer.
num_units_out: the number of output units in the layer.
activation: activation function.
stddev: the standard deviation for the weights.
bias: the initial value of the biases.
weight_decay: the weight decay.
batch_norm_params: parameters for the batch_norm. If is None don't use it.
is_training: whether or not the model is in training mode.
trainable: whether or not the variables should be trainable or not.
restore: whether or not the variables should be marked for restore.
scope: Optional scope for variable_op_scope.
Returns:
the tensor variable representing the result of the series of operations.
"""
with tf.variable_op_scope([inputs], scope, 'FC'):
num_units_in = inputs.get_shape()[1]
weights_shape = [num_units_in, num_units_out]
#weights_initializer = tf.truncated_normal_initializer(stddev=stddev)
weights_initializer = tf.random_normal_initializer(stddev=stddev)
l2_regularizer = lambda t: losses.l2_loss(t, weight_decay)
weights = variables.variable('weights',
shape=weights_shape,
initializer=weights_initializer,
regularizer=l2_regularizer,
trainable=trainable,
restore=restore)
if batch_norm_params is not None:
outputs = tf.matmul(inputs, weights)
with scopes.arg_scope([batch_norm],
is_training=is_training,
trainable=trainable,
restore=restore):
outputs = batch_norm(outputs, **batch_norm_params)
elif bias is None:
outputs = tf.matmul(inputs, weights)
else:
bias_shape = [
num_units_out,
]
bias_initializer = tf.constant_initializer(bias)
biases = variables.variable('biases',
shape=bias_shape,
initializer=bias_initializer,
trainable=trainable,
restore=restore)
outputs = tf.nn.xw_plus_b(inputs, weights, biases)
if activation:
outputs = activation(outputs)
return outputs
def conv2d(inputs,
num_filters_out,
kernel_size,
stride=1,
padding='SAME',
activation=tf.nn.relu,
stddev=0.01,
bias=None,
weight_decay=0,
batch_norm_params=None,
is_training=True,
trainable=True,
restore=True,
scope=None):
"""Adds a 2D convolution followed by an optional batch_norm layer.
conv2d creates a variable called 'weights', representing the convolutional
kernel, that is convolved with the input. If `batch_norm_params` is None, a
second variable called 'biases' is added to the result of the convolution
operation.
Args:
inputs: a tensor of size [batch_size, height, width, channels].
num_filters_out: the number of output filters.
kernel_size: a 2-D list comprising of the height and width of the filters.
stride: the stride in height and width of the convolution.
padding: one of 'VALID' or 'SAME'.
activation: activation function.
stddev: standard deviation of the truncated guassian weight distribution.
bias: the initial value of the biases.
weight_decay: the weight decay.
batch_norm_params: parameters for the batch_norm. If is None don't use it.
is_training: whether or not the model is in training mode.
trainable: whether or not the variables should be trainable or not.
restore: whether or not the variables should be marked for restore.
scope: Optional scope for variable_op_scope.
Returns:
a tensor representing the output of the operation.
Raises:
ValueError: if 'kernel_size' is not a 2-D list.
"""
if len(kernel_size) != 2:
raise ValueError('kernel_size must be a 2-D list.')
with tf.variable_op_scope([inputs], scope, 'Conv'):
num_filters_in = inputs.get_shape()[-1]
#num_filters_in = inputs.get_shape()[-3]
weights_shape = [
kernel_size[0], kernel_size[1], num_filters_in, num_filters_out
]
#weights_initializer = tf.truncated_normal_initializer(stddev=stddev)
stddev = math.sqrt(2. /
(kernel_size[0] * kernel_size[1] * num_filters_out))
weights_initializer = tf.random_normal_initializer(stddev=stddev)
l2_regularizer = lambda t: losses.l2_loss(t, weight_decay)
weights = variables.variable('weights',
shape=weights_shape,
initializer=weights_initializer,
regularizer=l2_regularizer,
trainable=trainable,
restore=restore)
conv = tf.nn.conv2d(inputs,
weights, [1, stride, stride, 1],
padding=padding)
#conv = tf.nn.conv2d(inputs, weights, [1, 1, stride, stride],
# padding=padding, data_format='NCHW')
if batch_norm_params is not None:
with scopes.arg_scope([batch_norm],
is_training=is_training,
trainable=trainable,
restore=restore):
outputs = batch_norm(conv, **batch_norm_params)
elif bias is None:
outputs = conv
else:
bias_shape = [
num_filters_out,
]
bias_initializer = tf.constant_initializer(bias)
biases = variables.variable('biases',
shape=bias_shape,
initializer=bias_initializer,
trainable=trainable,
restore=restore)
outputs = tf.nn.bias_add(conv, biases)
#outputs = tf.nn.bias_add(conv, biases, data_format='NCHW')
#outputs = tf.reshape(outputs, conv.get_shape())
if activation:
outputs = activation(outputs)
return outputs
def train_step_flow_extractor(args, params, use_flow_signal=False, supervised=False, n_points=None):
"""implements functions to perform ONE train step of the flow extractor"""
# initialize some recurrent parameters
n_points = params["n_points"] if n_points is None else n_points
use_deep_loss = not args.use_shallow_loss
device = params["device"]
# parse input for the loss module
# use_half_clouds = False if args.loss_type in ["triplet_l", "triplet_hinge", "js"] else True
select_inputs = select_inputs_(use_flow_signal, n_points, half_cloud=True)
out_acc_dict = defaultdict(lambda: 0.0)
out_error_dict = defaultdict(lambda: 0.0)
# initialize loss functions
if not supervised:
loss_func_shallow = initialize_flow_extractor_loss(loss_type=args.loss_type, params=params)
loss_func_deep = initialize_flow_extractor_loss(loss_type=params["loss_type_deep"], params=params)
deep_loss = deep_loss_(loss_func_deep, use_deep_loss, device, scale_type=args.deep_loss_scale)
else:
if args.train_type == 'knn':
if args.loss_type == 'knn':
loss_func_supervised = knn_loss
else:
loss_func_supervised = l2_loss(dim=1, norm=params["norm_FE"])
if (not supervised) and (params["cycle_consistency"] is not None):
cycons_func = cycons_func_(params["cycle_consistency"])
elif params["cycle_consistency_sup"] is not None:
cycons_func = cycons_func_(params["cycle_consistency_sup"])
def train_step_FE_func_uns(flow_extractor, cloud_embedder, c1, c2, flow_t1=None):
# train the flow extractor
flow_extractor.train()
cloud_embedder.eval()
flow_pred = flow_extractor(c1, c2)
c2_pred = c1 + flow_pred
c1_, c_anchor, c_negative, c_positive = select_inputs(c1, c2, c2_pred, flow_t1)
f_0, hidden_feats_0 = cloud_embedder(c1_, c_anchor)
f_p, hidden_feats_p = cloud_embedder(c1_, c_positive)
f_n, hidden_feats_n = cloud_embedder(c1_, c_negative)
loss_hidden_feats = deep_loss(hidden_feats_0, hidden_feats_p, hidden_feats_n)
loss_feat = loss_func_shallow(f_0, f_p, f_n)
loss_FE = loss_feat + loss_hidden_feats
out_error_dict["loss_feat"], out_error_dict["loss_hidden_feats"] = loss_feat.item(), loss_hidden_feats.item()
if params["cycle_consistency"] is not None:
c2_pred = c2_pred.detach()
if args.cycon_aug:
c2_pred = torch.cat((c2_pred, c2), dim=-1)
flow_pred_backwards = flow_extractor(c2_pred, c1)[..., :flow_pred.shape[-1]]
loss_cycons = cycons_func(flow_pred, flow_pred_backwards)
loss_FE = loss_cycons * args.cycon_contribution + loss_FE
out_error_dict["loss_cycons"] = loss_cycons.item()
c2_pred = c2_pred[..., :flow_pred.shape[-1]]
if params["local_consistency"]:
# loss_loccons = local_flow_consistency(pc1=c1, flow_pred=flow_pred)
loss_loccons = local_flow_consistency(c1, flow_pred)
loss_FE = loss_loccons * params["local_consistency"] + loss_FE
out_error_dict["loss_loccons"] = loss_loccons.item()
if params["chamfer"] > 0:
chamfer_dist = chamfer_distance(pc_pred=c2_pred, pc_target=c2)
loss_FE += params["chamfer"] * chamfer_dist
out_error_dict["chamfer"] = chamfer_dist.item()
if params["laplace"] > 0:
laplace_loss = laplacian_regularization(pc_pred=c2_pred, pc_target=c2)
loss_FE += params["laplace"] * laplace_loss
out_error_dict["laplace"] = laplace_loss.item()
if params["static_penalty"] > 0:
static_penalty = (c2_pred[..., :c1.shape[-1]] - c1).norm(dim=1).clamp(max=1).mean()
loss_FE += params["static_penalty"] * static_penalty
out_error_dict["static"] = static_penalty.item()
if flow_t1 is not None:
flow_errors_dict, flow_accs_dict = calculate_eucledian_metrics(flow_pred, flow_t1)
out_error_dict.update(flow_errors_dict)
out_acc_dict.update(flow_accs_dict)
out_dict = {"error": out_error_dict, "acc": out_acc_dict}
return loss_FE, out_dict
def train_step_FE_func_sup(flow_extractor, c1, c2, flow_t1=None):
# train the flow extractor
flow_extractor.train()
flow_pred = flow_extractor(c1, c2)
c2_pred = c1 + flow_pred
loss_FE = 0
# pdb.set_trace()
loss_FE += args.sup_scale * loss_func_supervised(flow_pred, flow_t1)
if params["cycle_consistency_sup"] is not None:
c2_pred_ = c2_pred.detach()
flow_pred_backwards = flow_extractor(c2_pred_, c1)
loss_cycons = cycons_func(flow_pred, flow_pred_backwards)
loss_FE = loss_cycons * args.cycon_contribution + loss_FE
out_error_dict["loss_cycons"] = loss_cycons.item()
if params["local_consistency"]:
loss_loccons = local_flow_consistency(c1, flow_pred)
loss_FE = loss_loccons * params["local_consistency"] + loss_FE
out_error_dict["loss_loccons"] = loss_loccons.item()
if params["chamfer"] > 0:
chamfer_dist = chamfer_distance(pc_pred=c2_pred, pc_target=c2)
loss_FE += params["chamfer"] * chamfer_dist
out_error_dict["chamfer"] = chamfer_dist.item()
if params["laplace"] > 0:
laplace_loss = laplacian_regularization(pc_pred=c2_pred, pc_target=c2)
loss_FE += params["laplace"] * laplace_loss
out_error_dict["laplace"] = laplace_loss.item()
if params["static_penalty"] > 0:
static_penalty = (c2_pred[..., :c1.shape[-1]] - c1).norm(dim=1).clamp(max=1).mean()
loss_FE += params["static_penalty"] * static_penalty
out_error_dict["static"] = static_penalty.item()
flow_errors_dict, flow_accs_dict = calculate_eucledian_metrics(flow_pred, flow_t1)
out_error_dict.update(flow_errors_dict)
out_acc_dict.update(flow_accs_dict)
out_dict = {"error": out_error_dict, "acc": out_acc_dict}
return loss_FE, out_dict
def train_step_FE_func_knn_sup(flow_extractor, c1, c2, flow_t1=None):
# train the flow extractor using knn loss
flow_extractor.train()
flow_pred = flow_extractor(c1, c2)
c2_pred = c1 + flow_pred
loss_FE = loss_func_supervised(c1 + flow_pred, c2)
if params["cycle_consistency_sup"] is not None:
c2_pred_ = c2_pred.detach()
flow_pred_backwards = flow_extractor(c2_pred_, c1)
loss_cycons = cycons_func(flow_pred, flow_pred_backwards)
loss_FE = loss_cycons * args.cycon_contribution + loss_FE
out_error_dict["loss_cycons"] = loss_cycons.item()
if params["local_consistency"]:
loss_loccons = local_flow_consistency(c1, flow_pred)
loss_FE = loss_loccons * params["local_consistency"] + loss_FE
out_error_dict["loss_loccons"] = loss_loccons.item()
if params["chamfer"] > 0:
chamfer_dist = chamfer_distance(pc_pred=c2_pred, pc_target=c2)
loss_FE += params["chamfer"] * chamfer_dist
out_error_dict["chamfer"] = chamfer_dist.item()
if params["laplace"] > 0:
laplace_loss = laplacian_regularization(pc_pred=c2_pred, pc_target=c2)
loss_FE += params["laplace"] * laplace_loss
out_error_dict["laplace"] = laplace_loss.item()
if params["static_penalty"] > 0:
static_penalty = (c2_pred[..., :c1.shape[-1]] - c1).norm(dim=1).clamp(max=1).mean()
loss_FE += params["static_penalty"] * static_penalty
out_error_dict["static"] = static_penalty.item()
flow_errors_dict, flow_accs_dict = calculate_eucledian_metrics(flow_pred, flow_t1)
out_error_dict.update(flow_errors_dict)
out_acc_dict.update(flow_accs_dict)
out_dict = {"error": out_error_dict, "acc": out_acc_dict}
return loss_FE, out_dict
if supervised:
if args.train_type == "knn":
return train_step_FE_func_knn_sup
return train_step_FE_func_sup
else:
return train_step_FE_func_uns
def train(self):
opt = self.opt
real_label = torch.ones(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
fake_label = torch.zeros(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
optimizer_de = optim.Adam(self.decoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_d = optim.Adam(self.netd.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_en = optim.Adam(self.encoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
l_bce = nn.BCELoss()
if opt.load_best_weights or opt.load_final_weights:
self.load_GAN_weights()
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
loss_de = []
loss_d = []
# train GAN
for epoch in range(opt.nepoch):
iternum = 0
lossd = 0
lossde = 0
self.decoder.train()
self.netd.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
inputdata = data[0].to(self.device)
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
z = z.unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(z)
# update netd
optimizer_d.zero_grad()
pred_real, feat_real = self.netd(inputdata)
pred_fake, feat_fake = self.netd(fake_image)
err_d_real = l_bce(pred_real, real_label)
err_d_fake = l_bce(pred_fake, fake_label)
err_d = (err_d_real + err_d_fake) * 0.5
err_d.backward()
optimizer_d.step()
# update netde
optimizer_de.zero_grad()
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
z = z.unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(z)
pred_fake, feat_fake = self.netd(fake_image)
err_g = l_bce(pred_fake, real_label)
err_g.backward()
optimizer_de.step()
# record loss
lossd += err_d.item()
lossde += err_g.item()
# record loss
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_d.append(lossd / opt.loss_iter)
loss_de.append(lossde / opt.loss_iter)
lossd = 0
lossde = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images')
visualize.save_images(train_img_dst, epoch, inputdata,
fake_image)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nepoch))
if opt.save_final_weight:
self.save_GAN_weights(opt.nepoch)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_d', loss_d)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_de', loss_de)
# train encoder
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
if opt.load_best_en_weights or opt.load_final_en_weights:
self.load_encoder_weights()
best_auc = 0
loss_en = []
self.netd.eval()
self.decoder.eval()
for epoch in range(opt.nenepoch):
iternum = 0
lossen = 0
self.encoder.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
inputdata = data[0].to(self.device)
# update encoder
optimizer_en.zero_grad()
latent_z = self.encoder(inputdata).squeeze()
latent_z = latent_z.unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_z)
pred_real, feat_real = self.netd(inputdata)
pred_fake, feat_fake = self.netd(fake_image)
k = 1.0
error_en = 1 / (inputdata.view(inputdata.shape[0], -1).shape[1]) * l2_loss(inputdata, fake_image) \
+ k * 1 / (feat_real.view(feat_real.shape[0], -1).shape[1]) * l2_loss(feat_real, feat_fake)
error_en.backward()
optimizer_en.step()
# record loss
lossen += error_en.item()
# record loss
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_en.append(lossen / opt.loss_iter)
lossen = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images2')
visualize.save_images(train_img_dst, epoch, inputdata,
fake_image)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nenepoch))
if epoch % opt.eva_epoch == 0:
performance = self.evaluate(epoch)
if performance['AUC'] > best_auc:
best_auc = performance['AUC']
if opt.save_best_en_weight:
self.save_encoder_weights(epoch, is_best=True)
# log performance
now = time.strftime("%c")
traintime = f'================ {now} ================\n'
visualize.write_to_log_file(
os.path.join(opt.outclass_dir, 'auc.log'), traintime)
visualize.log_current_performance(opt.outclass_dir,
performance, best_auc)
print(performance)
if opt.save_final_en_weight:
self.save_encoder_weights(opt.nepoch)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_en', loss_en)
print(">> Training model %s.[Done]" % self.name)
def train(self):
opt = self.opt
real_label = torch.ones(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
fake_label = torch.zeros(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
optimizer_f = optim.Adam(self.netf.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_de = optim.Adam(self.decoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_en = optim.Adam(self.encoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_d = optim.Adam(self.netd.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
l_bce = nn.BCELoss()
l_adv = l2_loss
l_con = nn.L1Loss()
l_enc = l2_loss
if opt.load_best_weights or opt.load_final_weights:
self.load_weights()
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
best_auc = 0
loss_eds = []
loss_fgs = []
loss_egs = []
for epoch in range(opt.nepoch):
iternum = 0
lossed = 0.0
lossfg = 0.0
losseg = 0.0
self.netf.train()
self.decoder.train()
self.encoder.train()
self.netd.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
data[0].requires_grad = True
inputdata = data[0].to(self.device)
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
latent_i = self.netf(z).unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_i)
latent_o = self.encoder(fake_image).squeeze()
real_latent = self.encoder(inputdata).squeeze()
pred_real = self.netd(real_latent)
pred_fake = self.netd(latent_o)
optimizer_en.zero_grad()
optimizer_d.zero_grad()
loss = loss_ed(pred_real, pred_fake)
# loss = losses.discriminator_logistic_simple_gp(pred_fake, pred_real, inputdata)
lossed += loss.item()
loss.backward()
optimizer_en.step()
optimizer_d.step()
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
latent_i = self.netf(z).unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_i)
latent_o = self.encoder(fake_image).squeeze()
pred_fake = self.netd(latent_o)
optimizer_f.zero_grad()
optimizer_de.zero_grad()
loss = loss_fg(pred_fake)
# loss = losses.generator_logistic_non_saturating(pred_fake)
lossfg += loss.item()
loss.backward()
optimizer_f.step()
optimizer_de.step()
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
latent_i = self.netf(z).unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_i)
latent_o = self.encoder(fake_image).squeeze()
optimizer_en.zero_grad()
optimizer_de.zero_grad()
loss = l2_loss(latent_i, latent_o)
# loss = losses.reconstruction(latent_i, latent_o)
losseg += loss.item()
loss.backward()
optimizer_en.step()
optimizer_de.step()
# record loss
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_eds.append(lossed / opt.loss_iter)
loss_fgs.append(lossfg / opt.loss_iter)
loss_egs.append(losseg / opt.loss_iter)
lossed = 0
lossfg = 0
losseg = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images')
visualize.save_images(train_img_dst, epoch, inputdata,
fake_image)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nepoch))
# if epoch % opt.eva_epoch == 0:
# performance = self.evaluate(epoch)
# if performance['AUC'] > best_auc:
# best_auc = performance['AUC']
# if opt.save_best_weight:
# self.save_weights(epoch, is_best=True)
# # log performance
# now = time.strftime("%c")
# traintime = f'================ {now} ================\n'
# visualize.write_to_log_file(os.path.join(opt.outclass_dir, 'auc.log'), traintime)
# visualize.log_current_performance(opt.outclass_dir, performance, best_auc)
# print(performance)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_ed', loss_eds)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_fg', loss_fgs)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_eg', loss_egs)
# if opt.save_final_weight:
# self.save_weights(opt.nepoch, is_best=False)
print(">> Training model %s.[Done]" % self.name)
