def _get_tgt_encoding_hand_code(self, observations, gps_pred,
compass_pred):
initial_pg = observations[self.goal_sensor_uuid]
pred_pg = self._update_pg(initial_pg, gps_pred, compass_pred)
if "pointgoal_with_gps" in observations:
true_pg = observations["pointgoal_with_gps"]
error = (pred_pg - true_pg).norm(dim=-1, keepdim=True)
if AuxLosses.is_active():
AuxLosses.register_loss("egomotion_error", error.mean(), 0.1)
valid_ego_preds = error.detach() < self.ego_error_threshold
pg = torch.where(valid_ego_preds, pred_pg.detach(), true_pg)
else:
pg = pred_pg
# Back-propping from the policy into the goal seems kinda odd -- doesn't make sense
# to move the goal to make the policy more/less likely to predict a given action
# We also have prefect supervision on what this should be!
goal_observations = _to_mag_and_unit_vec(pg.detach())
return self.tgt_embeding(goal_observations)
def forward(self, s, g):
if g is None:
return prior(s).sample()
priv_logits = self._priv(s, g)
d_cap_logits = self._d_cap(s)
d_cap = torch.sigmoid(d_cap_logits)
priv = torch.sigmoid(priv_logits)
x = d_cap * priv + (1 - d_cap) * prior(s).sample()
if AuxLosses.is_active():
priv_log_prob = F.logsigmoid(priv)
log_d_cap = F.logsigmoid(d_cap_logits)
AuxLosses.register_loss(
"information",
(
-d_cap * log_d_cap
+ (1 - d_cap) * (1 + priv_log_prob)
- (log_d_cap + priv_log_prob)
).mean(),
self.beta,
)
return x
def forward(self, s, obs):
if "pointgoal_with_gps" not in obs:
return self.prior(s).sample()
if self.use_odometry:
privileged_info = obs["pointgoal_with_gps_compass"]
else:
privileged_info = obs["pointgoal_with_gps"]
mu, sigma = torch.chunk(
self.encoder(
torch.cat(
[
s,
self.priv_embed(
_to_mag_and_unit_vec(privileged_info)
),
],
-1,
)
),
2,
s.dim() - 1,
)
if not self.use_info_bot:
mu.fill_(0.0)
sigma.fill_(1.0)
sigma = F.softplus(sigma)
dist = torch.distributions.Normal(mu, sigma)
x = dist.rsample()
# The following code block is for running with Selective Noise Injection
# if self.training:
# x = dist.rsample()
#else:
# x = dist.mean
if AuxLosses.is_active():
AuxLosses.register_loss(
"information",
_subsampled_mean(
torch.distributions.kl_divergence(dist, self.prior(s))
),
# torch.distributions.kl_divergence(dist, self.prior(s)).mean(),
self.beta,
)
return x
def evaluate_actions(
self, observations, prev_observations, rnn_hidden_states, prev_actions, masks, action
):
features, _ = self.net(observations, prev_observations, rnn_hidden_states, prev_actions, masks)
value = self.critic(features)
if self.supervise_stop:
stop_distribution = self.stop_action_distribution(features)
non_stop_distribution = self.non_stop_action_distribution(features)
action_log_probs = torch.where(
action == 0,
stop_distribution.log_probs(torch.full_like(action, 1)),
stop_distribution.log_probs(torch.full_like(action, 0))
+ non_stop_distribution.log_probs(
torch.max(action - 1, torch.zeros_like(action))
),
)
distribution_entropy = (
-1.0
* (
stop_distribution.probs[:, -1] * stop_distribution.logits[:, -1]
+ (
stop_distribution.probs[:, 0:1]
* non_stop_distribution.probs
* (
stop_distribution.logits[:, 0:1]
+ non_stop_distribution.logits
)
).sum(-1)
).mean()
)
stop_loss = F.cross_entropy(
stop_distribution.logits,
observations["stop_oracle"].long().squeeze(-1),
weight=torch.tensor(
[1.0, 1.0 / np.sqrt(100.0)], device=features.device
),
)
AuxLosses.register_loss("stop_loss", stop_loss)
else:
action_distribution = self.action_distribution(features)
action_log_probs = action_distribution.log_probs(action)
distribution_entropy = action_distribution.entropy().mean()
return value, action_log_probs, distribution_entropy
def forward(self, s, obs):
priv_emb = super().forward(s, obs)
gps = self.gps_head(s)
compass = self.compass_head(s)
pg = _update_pg_gps(obs["pointgoal"], gps)
embed_pg = self.predicted_embed(_to_mag_and_unit_vec(pg.detach()))
if AuxLosses.is_active():
AuxLosses.register_loss(
"egomotion_error",
torch.norm(pg - obs["pointgoal_with_gps"], dim=-1).mean(),
0.0,
)
AuxLosses.register_loss(
"compass_loss",
_subsampled_mean(_angular_distance_loss(compass, obs["compass"])),
)
AuxLosses.register_loss(
"gps_loss",
_subsampled_mean(
F.mse_loss(gps, obs["gps"], reduction="none").mean(-1)
),
)
return self.combine_layer(torch.cat([priv_emb, embed_pg], dim=-1))
def update(self, rollouts):
advantages = self.get_advantages(rollouts)
value_loss_epoch = 0
action_loss_epoch = 0
aux_losses_epoch = 0
dist_entropy_epoch = 0
AuxLosses.activate()
for e in range(self.ppo_epoch):
data_generator = rollouts.recurrent_generator(
advantages, self.num_mini_batch)
for sample in data_generator:
(
obs_batch,
prev_obs_batch,
recurrent_hidden_states_batch,
actions_batch,
prev_actions_batch,
value_preds_batch,
return_batch,
masks_batch,
old_action_log_probs_batch,
adv_targ,
) = sample
AuxLosses.clear()
# Reshape to do in a single forward pass for all steps
(
values,
action_log_probs,
dist_entropy,
) = self.actor_critic.evaluate_actions(
obs_batch,
prev_obs_batch,
recurrent_hidden_states_batch,
prev_actions_batch,
masks_batch,
actions_batch,
)
ratio = torch.exp(action_log_probs -
old_action_log_probs_batch)
surr1 = ratio * adv_targ
surr2 = (torch.clamp(ratio, 1.0 - self.clip_param,
1.0 + self.clip_param) * adv_targ)
action_loss = -torch.min(surr1, surr2).mean()
if self.use_clipped_value_loss:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(
-self.clip_param, self.clip_param)
value_losses = (values - return_batch).pow(2)
value_losses_clipped = (value_pred_clipped -
return_batch).pow(2)
value_loss = (
0.5 *
torch.max(value_losses, value_losses_clipped).mean())
else:
value_loss = 0.5 * (return_batch - values).pow(2).mean()
self.optimizer.zero_grad()
total_loss = (value_loss * self.value_loss_coef + action_loss -
dist_entropy * self.entropy_coef)
use_aux_loss = self.use_aux_losses
if use_aux_loss:
aux_losses = AuxLosses.reduce()
else:
aux_losses = AuxLosses.get_loss(
"information") * AuxLosses._loss_alphas["information"]
aux_losses_epoch += aux_losses.item()
total_loss = total_loss + aux_losses
self.before_backward(total_loss)
total_loss.backward()
self.after_backward(total_loss)
self.before_step()
self.optimizer.step()
self.after_step()
value_loss_epoch += value_loss.item()
action_loss_epoch += action_loss.item()
dist_entropy_epoch += dist_entropy.item()
num_updates = self.ppo_epoch * self.num_mini_batch
value_loss_epoch /= num_updates
action_loss_epoch /= num_updates
dist_entropy_epoch /= num_updates
aux_losses_epoch /= num_updates
AuxLosses.deactivate()
return (
value_loss_epoch,
action_loss_epoch,
dist_entropy_epoch,
aux_losses_epoch,
)
def forward(self, observations, prev_observations, rnn_hidden_states, prev_actions, masks):
if AuxLosses.is_active():
AuxLosses.obs = observations
depth_flag = False
rgb_flag = False
if "depth" in observations:
depth_flag = True
if "rgb" in observations:
rgb_flag = True
if "visual_features" in observations:
visual_features = observations["visual_features"]
prev_visual_features = observations["prev_visual_features"]
elif masks.size(0) != rnn_hidden_states.size(1):
obs_input = {}
N = rnn_hidden_states.size(1)
T = masks.size(0) // N
if depth_flag:
prev_obs = prev_observations["depth"].view(T, N, *prev_observations["depth"].size()[1:])
obs = observations["depth"].view(T, N, *observations["depth"].size()[1:])
obs_input["depth"] = torch.cat((prev_obs[0:1], obs), dim=0)
obs_input["depth"] = obs_input["depth"].view((T + 1) * N, *obs_input["depth"].size()[2:])
if rgb_flag:
prev_obs = prev_observations["rgb"].view(T, N, *prev_observations["rgb"].size()[1:])
obs = observations["rgb"].view(T, N, *observations["rgb"].size()[1:])
obs_input["rgb"] = torch.cat((prev_obs[0:1], obs), dim=0)
obs_input["rgb"] = obs_input["rgb"].view((T + 1) * N, *obs_input["rgb"].size()[2:])
obs_features = self.visual_encoder(obs_input)
prev_visual_features = obs_features[:T*N, :, :, :]
visual_features = obs_features[-T*N:, :, :, :]
else:
obs_input = {}
if depth_flag:
obs_input["depth"] = torch.cat((prev_observations["depth"], observations["depth"]), dim=0)
if rgb_flag:
obs_input["rgb"] = torch.cat((prev_observations["rgb"], observations["rgb"]), dim=0)
obs_features = self.visual_encoder(obs_input)
prev_visual_features, visual_features = obs_features.split(obs_features.size()[0] // 2, dim=0)
visual_features = self.compression(visual_features)
visual_emb = self.visual_fc(visual_features)
# difference of frames (unit 1)
flow_emb = self.visual_flow_encoder(
(visual_features - self.compression(prev_visual_features))
* masks.view(-1, 1, 1, 1)
)
prev_actions = self.prev_action_embedding(
((prev_actions.float() + 1) * masks).long().squeeze(-1)
)
context_emb = prev_actions + self._tgt_proj(observations["pointgoal"])
x, rnn_hidden_states = self.state_encoder(
torch.cat([visual_emb, flow_emb], dim=-1) + context_emb,
rnn_hidden_states,
masks,
)
tgt_encoding = self.get_tgt_encoding(observations, x)
x = torch.cat([x, tgt_encoding], dim=-1)
x = self.goal_mem_layer(x)
if AuxLosses.is_active():
n = rnn_hidden_states.size(1)
t = int(x.size(0) / n)
delta_ego = self.delta_egomotion_predictor(flow_emb).view(t, n, 3)
gps_gt = observations["gps"].view(t, n, 2)
compass_gt = observations["compass"].view(t, n, 1)
masks = masks.view(t, n, 1)
gt_delta = gps_gt[1:] - gps_gt[:-1]
gt_delta = _update_pg_gps_compass(
gt_delta.view((t - 1) * n, 2),
torch.zeros_like(gt_delta).view((t - 1) * n, 2),
compass_gt[:-1].view((t - 1) * n, 1),
).view(t - 1, n, 2)
AuxLosses.register_loss(
"delta_gps",
_subsampled_mean(
torch.masked_select(
F.mse_loss(
delta_ego[1:, :, 0:2], gt_delta, reduction="none"
).mean(dim=-1),
masks[1:, :, 0].bool(),
)
),
)
AuxLosses.register_loss(
"delta_compass",
_subsampled_mean(
torch.masked_select(
_angular_distance_loss(
delta_ego[1:, :, 2:], compass_gt[1:] - compass_gt[:-1]
),
masks[1:, :, 0].bool(),
)
),
)
return x, rnn_hidden_states
def train(self) -> None:
r"""Main method for DD-PPO.
Returns:
None
"""
self.local_rank, tcp_store = init_distrib_slurm(
self.config.RL.DDPPO.distrib_backend)
add_signal_handlers()
# Stores the number of workers that have finished their rollout
num_rollouts_done_store = distrib.PrefixStore("rollout_tracker",
tcp_store)
num_rollouts_done_store.set("num_done", "0")
self.world_rank = distrib.get_rank()
self.world_size = distrib.get_world_size()
self.config.defrost()
self.config.TORCH_GPU_ID = self.local_rank
self.config.SIMULATOR_GPU_ID = self.local_rank
self.config.TASK_CONFIG.TASK.POINTGOAL_WITH_EGO_PREDICTION_SENSOR.MODEL.GPU_ID = self.local_rank
# Multiply by the number of simulators to make sure they also get unique seeds
self.config.TASK_CONFIG.SEED += (self.world_rank *
self.config.NUM_PROCESSES)
self.config.freeze()
random.seed(self.config.TASK_CONFIG.SEED)
np.random.seed(self.config.TASK_CONFIG.SEED)
torch.manual_seed(self.config.TASK_CONFIG.SEED)
if torch.cuda.is_available():
self.device = torch.device("cuda", self.local_rank)
torch.cuda.set_device(self.device)
else:
self.device = torch.device("cpu")
self.envs = construct_envs(
self.config,
get_env_class(self.config.ENV_NAME),
workers_ignore_signals=True,
)
ppo_cfg = self.config.RL.PPO
if (not os.path.isdir(self.config.CHECKPOINT_FOLDER)
and self.world_rank == 0):
os.makedirs(self.config.CHECKPOINT_FOLDER)
self._setup_actor_critic_agent(ppo_cfg)
self.agent.init_distributed(find_unused_params=True)
if self.world_rank == 0:
logger.info("agent number of trainable parameters: {}".format(
sum(param.numel() for param in self.agent.parameters()
if param.requires_grad)))
observations = self.envs.reset()
batch = batch_obs(observations, device=self.device)
obs_space = self.envs.observation_spaces[0]
if self._static_encoder:
self._encoder = self.actor_critic.net.visual_encoder
obs_space = SpaceDict({
"visual_features":
spaces.Box(
low=np.finfo(np.float32).min,
high=np.finfo(np.float32).max,
shape=self._encoder.output_shape,
dtype=np.float32,
),
"prev_visual_features":
spaces.Box(
low=np.iinfo(np.uint32).min,
high=np.iinfo(np.uint32).max,
shape=self._encoder.output_shape,
dtype=np.float32,
),
**obs_space.spaces,
})
with torch.no_grad():
batch["visual_features"] = self._encoder(batch)
batch["prev_visual_features"] = torch.zeros_like(
batch["visual_features"])
rollouts = RolloutStorage(
ppo_cfg.num_steps,
self.envs.num_envs,
obs_space,
self.envs.action_spaces[0],
ppo_cfg.hidden_size,
num_recurrent_layers=self.actor_critic.net.num_recurrent_layers,
)
rollouts.to(self.device)
for sensor in rollouts.observations:
rollouts.observations[sensor][0].copy_(batch[sensor])
rollouts.previous_observations[sensor][0].copy_(
torch.zeros_like(batch[sensor]))
# batch and observations may contain shared PyTorch CUDA
# tensors. We must explicitly clear them here otherwise
# they will be kept in memory for the entire duration of training!
batch = None
observations = None
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
running_episode_stats = dict(
count=torch.zeros(self.envs.num_envs, 1, device=self.device),
reward=torch.zeros(self.envs.num_envs, 1, device=self.device),
)
window_episode_stats = defaultdict(
lambda: deque(maxlen=ppo_cfg.reward_window_size))
t_start = time.time()
env_time = 0
pth_time = 0
count_steps = 0
count_checkpoints = 0
start_update = 0
prev_time = 0
lr_scheduler = LambdaLR(
optimizer=self.agent.optimizer,
lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),
)
interrupted_state = load_interrupted_state()
if interrupted_state is not None:
self.agent.load_state_dict(interrupted_state["state_dict"])
self.agent.optimizer.load_state_dict(
interrupted_state["optim_state"])
lr_scheduler.load_state_dict(interrupted_state["lr_sched_state"])
requeue_stats = interrupted_state["requeue_stats"]
env_time = requeue_stats["env_time"]
pth_time = requeue_stats["pth_time"]
count_steps = requeue_stats["count_steps"]
count_checkpoints = requeue_stats["count_checkpoints"]
start_update = requeue_stats["start_update"]
prev_time = requeue_stats["prev_time"]
with (TensorboardWriter(self.config.TENSORBOARD_DIR,
flush_secs=self.flush_secs)
if self.world_rank == 0 else contextlib.suppress()) as writer:
for update in range(start_update, self.config.NUM_UPDATES):
if ppo_cfg.use_linear_lr_decay:
lr_scheduler.step()
self.actor_critic.update_ib_beta(count_steps)
if ppo_cfg.use_linear_clip_decay:
self.agent.clip_param = ppo_cfg.clip_param * linear_decay(
update, self.config.NUM_UPDATES)
if EXIT.is_set():
self.envs.close()
if REQUEUE.is_set() and self.world_rank == 0:
requeue_stats = dict(
env_time=env_time,
pth_time=pth_time,
count_steps=count_steps,
count_checkpoints=count_checkpoints,
start_update=update,
prev_time=(time.time() - t_start) + prev_time,
)
save_interrupted_state(
dict(
state_dict=self.agent.state_dict(),
optim_state=self.agent.optimizer.state_dict(),
lr_sched_state=lr_scheduler.state_dict(),
config=self.config,
requeue_stats=requeue_stats,
))
requeue_job()
return
count_steps_delta = 0
self.agent.eval()
for step in range(ppo_cfg.num_steps):
(
delta_pth_time,
delta_env_time,
delta_steps,
) = self._collect_rollout_step(rollouts,
current_episode_reward,
running_episode_stats)
pth_time += delta_pth_time
env_time += delta_env_time
count_steps_delta += delta_steps
# This is where the preemption of workers happens. If a
# worker detects it will be a straggler, it preempts itself!
if (step >=
ppo_cfg.num_steps * self.SHORT_ROLLOUT_THRESHOLD
) and int(num_rollouts_done_store.get("num_done")) > (
self.config.RL.DDPPO.sync_frac * self.world_size):
break
num_rollouts_done_store.add("num_done", 1)
self.agent.train()
if self._static_encoder:
self._encoder.eval()
(
delta_pth_time,
value_loss,
action_loss,
dist_entropy,
aux_loss,
) = self._update_agent(ppo_cfg, rollouts)
pth_time += delta_pth_time
stats_ordering = list(sorted(running_episode_stats.keys()))
stats = torch.stack(
[running_episode_stats[k] for k in stats_ordering], 0)
distrib.all_reduce(stats)
for i, k in enumerate(stats_ordering):
window_episode_stats[k].append(stats[i].clone())
stats = torch.tensor(
[
value_loss,
action_loss,
aux_loss,
AuxLosses.get_loss("egomotion_error"),
AuxLosses.get_loss("information"),
0.1 * dist_entropy,
count_steps_delta,
],
device=self.device,
)
distrib.all_reduce(stats)
count_steps += stats[-1].item()
if self.world_rank == 0:
num_rollouts_done_store.set("num_done", "0")
losses = [
stats[i].item() / self.world_size
for i in range(stats.size(0) - 1)
]
deltas = {
k: ((v[-1] - v[0]).sum().item()
if len(v) > 1 else v[0].sum().item())
for k, v in window_episode_stats.items()
}
deltas["count"] = max(deltas["count"], 1.0)
print(deltas["reward"])
writer.add_scalar(
"reward",
deltas["reward"] / deltas["count"],
count_steps,
)
# Check to see if there are any metrics
# that haven't been logged yet
metrics = {
k: v / deltas["count"]
for k, v in deltas.items()
if k not in {"reward", "count"}
}
if len(metrics) > 0:
writer.add_scalars("metrics", metrics, count_steps)
writer.add_scalars(
"losses",
{
k: l
for l, k in zip(losses, [
"value", "policy", "aux", "egomotion_error",
"information", "entropy"
])
},
count_steps,
)
# log stats
if update > 0 and update % self.config.LOG_INTERVAL == 0:
logger.info("update: {}\tfps: {:.3f}\t".format(
update,
count_steps /
((time.time() - t_start) + prev_time),
))
logger.info(
"update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
"frames: {}".format(update, env_time, pth_time,
count_steps))
logger.info("Average window size: {} {}".format(
len(window_episode_stats["count"]),
" ".join(
"{}: {:.3f}".format(k, v / deltas["count"])
for k, v in deltas.items() if k != "count"),
))
# checkpoint model
if update % self.config.CHECKPOINT_INTERVAL == 0:
self.save_checkpoint(
f"ckpt.{count_checkpoints}.pth",
dict(step=count_steps),
)
requeue_stats = dict(
env_time=env_time,
pth_time=pth_time,
count_steps=count_steps,
count_checkpoints=count_checkpoints,
start_update=update,
prev_time=(time.time() - t_start) + prev_time,
)
save_interrupted_state(
dict(
state_dict=self.agent.state_dict(),
optim_state=self.agent.optimizer.state_dict(),
lr_sched_state=lr_scheduler.state_dict(),
config=self.config,
requeue_stats=requeue_stats,
))
count_checkpoints += 1
self.envs.close()
def forward(self, observations, rnn_hidden_states, prev_actions, masks):
if AuxLosses.is_active():
AuxLosses.obs = observations
if "visual_features" in observations:
visual_features = observations["visual_features"]
else:
visual_features = self.visual_encoder(observations)
visual_features = self.compression(visual_features)
visual_emb = self.visual_fc(visual_features)
flow_emb = self.visual_flow_encoder(
(visual_features - self.compression(observations["prev_visual_features"]))
* masks.view(-1, 1, 1, 1)
)
prev_actions = self.prev_action_embedding(
((prev_actions.float() + 1) * masks).long().squeeze(-1)
)
context_emb = prev_actions + self._tgt_proj(observations["pointgoal"])
x, rnn_hidden_states = self.state_encoder(
torch.cat([visual_emb, flow_emb], dim=-1) + context_emb,
rnn_hidden_states,
masks,
)
tgt_encoding = self.get_tgt_encoding(observations, x)
x = torch.cat([x, tgt_encoding], dim=-1)
x = self.goal_mem_layer(x)
if AuxLosses.is_active():
n = rnn_hidden_states.size(1)
t = int(x.size(0) / n)
delta_ego = self.delta_egomotion_predictor(flow_emb).view(t, n, 3)
gps_gt = observations["gps"].view(t, n, 2)
compass_gt = observations["compass"].view(t, n, 1)
masks = masks.view(t, n, 1)
gt_delta = gps_gt[1:] - gps_gt[:-1]
gt_delta = _update_pg_gps_compass(
gt_delta.view((t - 1) * n, 2),
torch.zeros_like(gt_delta).view((t - 1) * n, 2),
compass_gt[:-1].view((t - 1) * n, 1),
).view(t - 1, n, 2)
AuxLosses.register_loss(
"delta_gps",
_subsampled_mean(
torch.masked_select(
F.mse_loss(
delta_ego[1:, :, 0:2], gt_delta, reduction="none"
).mean(dim=-1),
masks[1:, :, 0].bool(),
)
),
)
AuxLosses.register_loss(
"delta_compass",
_subsampled_mean(
torch.masked_select(
_angular_distance_loss(
delta_ego[1:, :, 2:], compass_gt[1:] - compass_gt[:-1]
),
masks[1:, :, 0].bool(),
)
),
)
return x, rnn_hidden_states
def seq_forward(self, x, context, mems, masks):
r"""Forward for a sequence of length T
Args:
x: (T, N, -1) Tensor that has been flattened to (T * N, -1)
hidden_states: The starting hidden state.
masks: The masks to be applied to hidden state at every timestep.
A (T, N) tensor flatten to (T * N)
"""
# x is a (T, N, -1) tensor flattened to (T * N, -1)
n = mems.size(1)
t = int(x.size(0) / n)
# unflatten
x = x.view(t, n, x.size(1))
context = context.view(t, n, context.size(1))
masks = masks.view(t, n)
if self.self_sup:
ep_lens = []
for i in range(n):
last_zero = 0
has_zeros = ((
masks[1:-1,
i] == 0.0).nonzero().squeeze(-1).cpu().unbind(0))
for z in has_zeros:
z = z.item() + 1
ep_lens.append(z - last_zero)
last_zero = z
ep_lens.append(t - last_zero)
k = random.randint(
1,
max(
min(self.max_self_sup_K,
int(0.8 * np.mean(np.array(ep_lens)))), 2),
)
else:
k = None
content, mems, query = self.transformer.transformer_seq_forward(
x, context, mems, masks, two_stream_k=k)
if self.self_sup:
positives = x
negatives = []
for _ in range(3):
negative_inds = torch.randperm(t * n, device=x.device)
negatives.append(
torch.gather(
x.view(t * n, -1),
dim=0,
index=negative_inds.view(t * n,
1).expand(t * n, x.size(-1)),
).view(t, n, -1))
negatives = torch.stack(negatives, dim=-1)
positives = torch.einsum("...i, ...i -> ...", positives, query)
negatives = torch.einsum("...ik, ...i -> ...k", negatives, query)
cpc_logits = torch.stack([positives.unsqueeze(-1), negatives],
dim=-1)
valid_modeling_queries = torch.ones(t,
n,
device=query.device,
dtype=torch.bool)
valid_modeling_queries[0:k] = 0
for i in range(n):
has_zeros_batch = ((
masks[:, i] == 0.0).nonzero().squeeze(-1).cpu().unbind(0))
for z in has_zeros_batch:
valid_modeling_queries[n:n + k, i] = 0
cpc_loss = torch.masked_select(
F.cross_entropy(
cpc_logits,
torch.zeros(t,
n,
dtype=torch.long,
device=cpc_logits.device),
reduction="none",
),
valid_modeling_queries,
).mean()
AuxLosses.register_loss("CPC|A", cpc_loss, 0.2)
content = content.view(t * n, -1)
return content, mems