lapjv和linear_sum_assignment的返回值的区别

cost, row_ind, col_ind = lapjv(cost_matrix)

print(f"行分配: {row_ind}") # 每一行映射到的最优列

print(f"列分配: {col_ind}") # 每一列映射到的最优行

所以,row_ind[i]是第i行对应的列索引,而 col_ind[j]是第j列对应的行索引。

对于linear_sum_assignment的row_ind和col_ind,根据scipy的文档,返回的row_ind和col_ind是行和列的索引对

def example_compare2():


    a = np.array([
        [1.46439595e+03, 2.69431832e+03, 1.29987563e-01, 1.55361073e+03, 1.44865847e+03, 3.10717022e+03,
         2.57535292e+03],
        [2.57436453e+03, 3.10716454e+03, 1.44876144e+03, 1.55268757e+03, 2.47319823e-02, 2.69476280e+03,
         1.46539683e+03],
        [1.76492370e-01, 1.64718736e+03, 1.46437509e+03, 1.55696505e+03, 2.57452162e+03, 2.79781559e+03,
         3.09548074e+03],
        [3.09534496e+03, 2.79864587e+03, 2.57524914e+03, 1.55557721e+03, 1.46499350e+03, 1.64802488e+03,
         4.88675400e-01],
        [1.55677233e+03, 1.55718485e+03, 1.55318774e+03, 5.25698984e-01, 1.55220347e+03, 1.55606378e+03,
         1.55539410e+03],
        [1.64741113e+03, 3.21426508e-01, 2.69453589e+03, 1.55685445e+03, 3.10730844e+03, 1.65236984e+03,
         2.79834747e+03],
        [2.79802807e+03, 1.65309575e+03, 3.10731726e+03, 1.55578018e+03, 2.69459113e+03, 5.78065695e-01, 1.64716845e+03]
    ])

    # 使用 lapjv
    c, row_ind_lapjv, col_ind_lapjv = lapjv(a)
    lapjv_total_cost = a[np.arange(len(row_ind_lapjv)), row_ind_lapjv].sum()
    print("lapjv分配结果:")
    print(f"行分配(row_ind): {row_ind_lapjv}")
    print(f"列分配(col_ind): {col_ind_lapjv}")
    print(f"总代价: {lapjv_total_cost}")

    # 使用 linear_sum_assignment
    row_ind_linear, col_ind_linear = linear_sum_assignment(a)
    linear_sum_assignment_total_cost = a[row_ind_linear, col_ind_linear].sum()
    print("\nlinear_sum_assignment分配结果:")
    print(f"行分配(row_ind): {row_ind_linear}")
    print(f"列分配(col_ind): {col_ind_linear}")
    print(f"总代价: {linear_sum_assignment_total_cost}")

lapjv分配结果:
行分配(row_ind): [2 4 0 6 3 1 5]
列分配(col_ind): [2 5 0 4 1 6 3]
总代价: 2.2450785023

linear_sum_assignment分配结果:
行分配(row_ind): [0 1 2 3 4 5 6]
列分配(col_ind): [2 4 0 6 3 1 5]
总代价: 2.2450785023

遇到的问题:

用linear_sum_assignment时出现了gui退出报错的问题,原因暂时未知,而且错误不是每次都报,不好解决。所以用lapjv代替。

Process finished with exit code -1073741819 (0xC0000005) 

from typing import Dict, Tuple import torch import torch.nn.functional as F import numpy as np from einops import rearrange from diffusers.schedulers.scheduling_ddpm import DDPMScheduler import inspect from controller.model.common.normalizer import LinearNormalizer from controller.policy.base_image_policy import BaseImagePolicy from controller.model.diffusion.transformer_for_action_diffusion import TransformerForActionDiffusion from controller.model.vision.obs_encoder import ObsEncoder from scipy.optimize import linear_sum_assignment import pickle # Adapted from https://github.com/lucidrains/pi-zero-pytorch/blob/e82fced40e55023a0ded22ab3bda495964353253/pi_zero_pytorch/pi_zero.py#L216 def noise_assignment(data, noise): device = data.device data, noise = tuple(rearrange(t, 'b ... -> b (...)') for t in (data, noise)) dist = torch.cdist(data, noise) _, assign = linear_sum_assignment(dist.cpu()) return torch.from_numpy(assign).to(device) class DexGraspVLAController(BaseImagePolicy): def __init__(self, shape_meta: dict, noise_scheduler: DDPMScheduler, obs_encoder: ObsEncoder, num_inference_steps=None, # arch n_layer=7, n_head=8, p_drop_attn=0.1, use_attn_mask=False, start_ckpt_path=None, # parameters passed to step **kwargs): super().__init__() # parse shapes action_shape = shape_meta['action']['shape'] assert len(action_shape) == 1 action_dim = action_shape[0] action_horizon = shape_meta['action']['horizon'] obs_shape, obs_part_length = obs_encoder.output_shape() n_emb = obs_shape[-1] obs_tokens = obs_shape[-2] model = TransformerForActionDiffusion( input_dim=action_dim, output_dim=action_dim, action_horizon=action_horizon, n_layer=n_layer, n_head=n_head, n_emb=n_emb, max_cond_tokens=obs_tokens+1, # obs tokens + 1 token for time p_drop_attn=p_drop_attn, obs_part_length=obs_part_length, use_attn_mask=use_attn_mask ) self.obs_encoder = obs_encoder self.model = model self.noise_scheduler = noise_scheduler self.normalizer = LinearNormalizer() self.action_dim = action_dim self.action_horizon = action_horizon self.start_ckpt_path = start_ckpt_path self.kwargs = kwargs if num_inference_steps is None: num_inference_steps = noise_scheduler.config.num_train_timesteps self.num_inference_steps = num_inference_steps # ========= inference ============ def conditional_sample(self, cond=None, gen_attn_map=True, **kwargs): model = self.model scheduler = self.noise_scheduler B = cond.shape[0] trajectory = torch.randn( size=(B, self.action_horizon, self.action_dim), dtype=self.dtype, device=self.device) # set step values scheduler.set_timesteps(self.num_inference_steps) # Store attention maps for all timesteps all_timestep_attention_maps = {} for t in scheduler.timesteps: # 1. predict model output model_output, attention_maps = model(trajectory, t, cond, training=False, gen_attn_map=gen_attn_map) all_timestep_attention_maps[t.cpu().item()] = attention_maps # 2. compute previous image: x_t -> x_t-1 trajectory = scheduler.step( model_output, t, trajectory, **kwargs ).prev_sample return trajectory, all_timestep_attention_maps def predict_action(self, obs_dict: Dict[str, torch.Tensor], output_path: str = None) -> Dict[str, torch.Tensor]: """ obs_dict: must include "obs" key action_pred: predicted action """ assert 'past_action' not in obs_dict # not implemented yet # normalize input # nobs = self.normalizer.normalize(obs_dict) nobs = obs_dict B = next(iter(nobs.values())).shape[0] # process input obs_tokens = self.obs_encoder(nobs, training=False) # (B, N, n_emb) # run sampling nsample, all_timestep_attention_maps = self.conditional_sample( cond=obs_tokens, gen_attn_map=True if output_path is not None else False, **self.kwargs) # unnormalize prediction assert nsample.shape == (B, self.action_horizon, self.action_dim) action_pred = self.normalizer['action'].unnormalize(nsample) if output_path is not None: # Convert tensors in obs_dict to numpy arrays obs_dict_numpy = {} for k, v in obs_dict.items(): if k in ['rgbm', 'right_cam_img']: obs_dict_numpy[k] = np.clip(v.detach().cpu().numpy() * 255, 0, 255).astype(np.uint8) else: obs_dict_numpy[k] = v.detach().cpu().numpy() obs_dict_numpy[k] = obs_dict_numpy[k][:2] save_dict = { 'attention_maps': all_timestep_attention_maps, 'obs_dict': obs_dict_numpy } with open(output_path, 'wb') as f: pickle.dump(save_dict, f) return action_pred # ========= training ============ def set_normalizer(self, normalizer: LinearNormalizer): self.normalizer.load_state_dict(normalizer.state_dict()) def get_optimizer( self, lr: float, weight_decay: float, betas: Tuple[float, float], ) -> torch.optim.Optimizer: # start with all of the candidate parameters (that require grad) param_dict = {pn: p for pn, p in self.named_parameters()} param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] optim_groups = [ {'params': decay_params, 'weight_decay': weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0} ] num_decay_params = sum(p.numel() for p in decay_params) num_nodecay_params = sum(p.numel() for p in nodecay_params) print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters print(f"Fused AdamW available: {fused_available}") optimizer = torch.optim.AdamW( optim_groups, lr=lr, betas=betas, fused=fused_available ) return optimizer def compute_loss(self, batch, training=True): # normalize input assert 'valid_mask' not in batch # nobs = self.normalizer.normalize(batch['obs']) nobs = batch['obs'] nactions = self.normalizer['action'].normalize(batch['action']) trajectory = nactions # process input obs_tokens = self.obs_encoder(nobs, training) # (B, N, n_emb) # Sample noise that we'll add to the images noise = torch.randn(trajectory.shape, device=trajectory.device) assignment = noise_assignment(trajectory, noise) noise = noise[assignment] # Sample a random timestep for each image timesteps = torch.randint( 0, self.noise_scheduler.config.num_train_timesteps, (nactions.shape[0],), device=trajectory.device ).long() # Add noise to the clean images according to the noise magnitude at each timestep # (this is the forward diffusion process) noisy_trajectory = self.noise_scheduler.add_noise( trajectory, noise, timesteps) # Predict the noise residual pred, _ = self.model( noisy_trajectory, timesteps, cond=obs_tokens, training=training, gen_attn_map=False ) pred_type = self.noise_scheduler.config.prediction_type if pred_type == 'epsilon': target = noise elif pred_type == 'sample': target = trajectory else: raise ValueError(f"Unsupported prediction type {pred_type}") loss = F.mse_loss(pred, target) return loss def forward(self, batch, training=True): return self.compute_loss(batch, training)全文注释
07-16
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值