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)