def single_frame_loss(): JOB = "20210124114528" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" device = torch.device("cuda:0") # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) ( train_dataset, train_loader, valid_dataset, valid_loader, test_dataset, test_loader, ) = init_datasets(opts) model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model = model.to(device) model.eval() criterion = nn.CrossEntropyLoss() test_loss_one = 0.0 test_loss_all = 0.0 n_test = 0 with torch.no_grad(): for test_tensors in test_loader: # Skip bad sequences. if len(test_tensors["player_idxs"]) < model.seq_len: continue player_trajs = test_tensors["player_trajs"].flatten() n_player_trajs = len(player_trajs) labels = player_trajs.to(device) preds = model(test_tensors)["player"][:n_player_trajs] test_loss_one += criterion(preds[:10], labels[:10]).item() test_loss_all += criterion(preds, labels).item() n_test += 1 print(test_loss_one / n_test) print(test_loss_all / n_test)
def plot_attn_through_time(): JOB = "20210124114950" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) (train_dataset, _, _, _, test_dataset, _) = init_datasets(opts) model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model.eval() court = plt.imread("court.png") width = 5 height = width * court.shape[0] / float(court.shape[1]) pt_scale = 1.4 home_dir = os.path.expanduser("~") os.makedirs(f"{home_dir}/test", exist_ok=True) (fig, ax) = plt.subplots(figsize=(width, height)) save_output = SaveOutput() hook_handles = [] for layer in model.transformer.layers: hook_handles.append(layer.self_attn.register_forward_hook(save_output)) test_idx = 851 tensors = test_dataset[test_idx] with torch.no_grad(): _ = model(tensors) colors = [] markers = [] for p_idx in range(10): if tensors["player_hoop_sides"][0, p_idx]: colors.append("white") markers.append("s") else: colors.append("gray") markers.append("^") ball_start = 10 * len(tensors["player_idxs"]) for time_step in range(19): print(time_step, flush=True) ball_pos = ball_start + time_step ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN]) ax.axis("off") ax.grid(False) for p_idx in range(10): ax.scatter( [tensors["player_xs"][time_step, p_idx]], [tensors["player_ys"][time_step, p_idx]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c=colors[p_idx], marker=markers[p_idx], edgecolors="black", ) ax.scatter( [tensors["ball_xs"][time_step]], [tensors["ball_ys"][time_step]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c="#ff8c00", edgecolors="black", ) plt.subplots_adjust(0, 0, 1, 1) fig.canvas.draw() court_img = Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb()) plt.cla() time_step = str(time_step).zfill(2) court_img.save(f"{home_dir}/test/court_{time_step}.png") for attn_layer_idx in range(len(save_output.outputs)): attn_scores = (save_output.outputs[attn_layer_idx][1] [0].detach().cpu().numpy()) for attn_head_idx in range(len(attn_scores)): attn_head_scores = attn_scores[attn_head_idx, ball_pos] player_scores = (attn_head_scores[:ball_start].reshape( -1, 10).sum(axis=0)) vmax = player_scores.max() player_scores_normed = player_scores / vmax player_scores_normed = np.nan_to_num(player_scores_normed, 0) player_colors = [ str(player_score) for player_score in player_scores_normed ] ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN]) ax.axis("off") ax.grid(False) for p_idx in range(10): ax.scatter( [tensors["player_xs"][time_step, p_idx]], [tensors["player_ys"][time_step, p_idx]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c=player_colors[p_idx], marker=markers[p_idx], cmap="gray", vmin=0, vmax=1, edgecolors="black", ) ball_score = attn_head_scores[ball_start:].sum() ax.scatter( [tensors["ball_xs"][time_step]], [tensors["ball_ys"][time_step]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c="#ff8c00", edgecolors="black", ) plt.subplots_adjust(0, 0, 1, 1) fig.canvas.draw() court_img = Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb()) plt.cla() ball_score = str(ball_score).split(".")[1][:2] img_id = f"{str(attn_layer_idx).zfill(2)}_{str(attn_head_idx).zfill(2)}" time_step = str(time_step).zfill(2) court_img.save( f"{home_dir}/test/{img_id}_{time_step}_{ball_score}.png") shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test") shutil.rmtree(f"{home_dir}/test")
def naive_player_test_set_perplexity(): JOB = "20210124114528" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" device = torch.device("cuda:0") # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) ( train_dataset, train_loader, valid_dataset, valid_loader, test_dataset, test_loader, ) = init_datasets(opts) default_hz = 25 hz = 5 skip = default_hz // hz skip_secs = 1 / skip all_player_trajs = [] for gameid in train_dataset.gameids: X = np.load(f"{GAMES_DIR}/{gameid}_X.npy") seq_data = X[::skip] player_xs = seq_data[:, 20:30] player_ys = seq_data[:, 30:40] player_x_diffs = np.diff(player_xs, axis=0) player_y_diffs = np.diff(player_ys, axis=0) wall_clock_diffs = np.diff(seq_data[:, -1]) / 1000 keep_diffs = wall_clock_diffs <= 1.5 * skip_secs player_x_diffs = player_x_diffs[keep_diffs].flatten() player_y_diffs = player_y_diffs[keep_diffs].flatten() player_traj_rows = np.digitize(player_y_diffs, train_dataset.player_traj_bins) player_traj_cols = np.digitize(player_x_diffs, train_dataset.player_traj_bins) player_trajs = player_traj_rows * train_dataset.player_traj_n + player_traj_cols all_player_trajs.append(player_trajs) all_player_trajs = np.concatenate(all_player_trajs) (unique, counts) = np.unique(all_player_trajs, return_counts=True) d_counts = dict(zip(unique, counts)) probs = np.zeros(train_dataset.ball_traj_n**3) for (traj, count) in d_counts.items(): probs[traj] = count probs = probs / probs.sum() model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model = model.to(device) model.eval() test_loss = 0.0 n_test = 0 with torch.no_grad(): for test_tensors in test_loader: # Skip bad sequences. if len(test_tensors["player_idxs"]) < model.seq_len: continue labels = test_tensors["player_trajs"].flatten() test_loss -= np.log(probs[labels]).mean() n_test += 1 print(test_loss) print(np.exp(test_loss / n_test))
def plot_generated_trajectories(): JOB = "20210124114528" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) (train_dataset, _, _, _, test_dataset, _) = init_datasets(opts) model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model.eval() cand_test_idxs = [] for test_idx in range(len(test_dataset.gameids)): tensors = test_dataset[test_idx] if len(tensors["player_idxs"]) == model.seq_len: cand_test_idxs.append(test_idx) court = plt.imread("court.png") width = 5 height = width * court.shape[0] / float(court.shape[1]) pt_scale = 1.4 player_traj_bins = np.array(list(test_dataset.player_traj_bins) + [5.5]) home_dir = os.path.expanduser("~") os.makedirs(f"{home_dir}/test", exist_ok=True) (fig, ax) = plt.subplots(figsize=(width, height)) saved = 0 np.random.seed(2010) while saved < 20: test_idx = np.random.choice(cand_test_idxs) cand_test_idxs.remove(test_idx) print(test_idx, flush=True) for start_time_step in [0, 5, 9, 19]: tensors = test_dataset[test_idx] (player_idxs, p_idxs) = tensors["player_idxs"][0].sort() player_traj_n = test_dataset.player_traj_n seq_len = len(tensors["player_idxs"]) grid_gap = np.diff(test_dataset.player_traj_bins)[-1] / 2 with torch.no_grad(): for step in range(seq_len - start_time_step - 1): preds_start = 10 * (start_time_step + step) preds_stop = preds_start + 10 preds = model(tensors)["player"][preds_start:preds_stop] probs = torch.softmax(preds, dim=1) samp_trajs = torch.multinomial(probs, 1) samp_rows = samp_trajs // player_traj_n samp_cols = samp_trajs % player_traj_n samp_xs = (player_traj_bins[samp_cols] - grid_gap + np.random.uniform(-grid_gap, grid_gap)) samp_ys = (player_traj_bins[samp_rows] - grid_gap + np.random.uniform(-grid_gap, grid_gap)) tensors["player_xs"][start_time_step + step + 1] = ( tensors["player_xs"][start_time_step + step] + samp_xs.flatten()) tensors["player_ys"][start_time_step + step + 1] = ( tensors["player_ys"][start_time_step + step] + samp_ys.flatten()) markers = [] for p_idx in range(10): if tensors["player_hoop_sides"][0, p_idx]: markers.append("s") else: markers.append("^") ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN]) ax.axis("off") ax.grid(False) for (idx, p_idx) in enumerate(p_idxs): ax.scatter( [tensors["player_xs"][0, p_idx]], [tensors["player_ys"][0, p_idx]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c=colors[idx], marker=markers[p_idx], edgecolors="black", ) plt.xlim(auto=False) plt.ylim(auto=False) ax.plot( tensors["player_xs"][:start_time_step + 1, p_idx], tensors["player_ys"][:start_time_step + 1, p_idx], c=colors[idx], ) ax.plot( tensors["player_xs"][start_time_step:, p_idx], tensors["player_ys"][start_time_step:, p_idx], c=colors[idx], linestyle=":", ) ax.plot( tensors["ball_xs"], tensors["ball_ys"], c="#ff8c00", ) ax.scatter( [tensors["ball_xs"][0]], [tensors["ball_ys"][0]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c="#ff8c00", edgecolors="black", ) plt.subplots_adjust(0, 0, 1, 1) fig.canvas.draw() court_img = Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb()) plt.cla() court_img.save( f"{home_dir}/test/trajectories_{test_idx}_{start_time_step}.png" ) saved += 1 shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test") shutil.rmtree(f"{home_dir}/test")
def plot_traj_preds(): JOB = "20210124114528" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) (train_dataset, _, _, _, test_dataset, _) = init_datasets(opts) model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model.eval() cand_test_idxs = [] for test_idx in range(len(test_dataset.gameids)): tensors = test_dataset[test_idx] if len(tensors["player_idxs"]) == model.seq_len: cand_test_idxs.append(test_idx) court = plt.imread("court.png") width = 5 height = width * court.shape[0] / float(court.shape[1]) vert_buffer = 10 pt_scale = 1.4 scale = 8 pred_player_idx = 5 home_dir = os.path.expanduser("~") os.makedirs(f"{home_dir}/test", exist_ok=True) (fig, ax) = plt.subplots(figsize=(width, height)) saved = 0 np.random.seed(2010) while saved < 20: test_idx = np.random.choice(cand_test_idxs) cand_test_idxs.remove(test_idx) print(test_idx, flush=True) tensors = test_dataset[test_idx] player_traj_n = test_dataset.player_traj_n with torch.no_grad(): seq_len = len(tensors["player_idxs"]) labels = tensors["player_trajs"].flatten() preds = model(tensors)["player"][:len(labels)] probs = torch.softmax(preds, dim=1) pred_player_probs = probs.view(seq_len, -1, player_traj_n**2)[:, pred_player_idx] colors = [] markers = [] for p_idx in range(10): if p_idx == pred_player_idx: colors.append("r") elif tensors["player_hoop_sides"][0, p_idx]: colors.append("white") else: colors.append("gray") if tensors["player_hoop_sides"][0, p_idx]: markers.append("s") else: markers.append("^") for time_step in range(len(tensors["player_idxs"])): ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN]) ax.axis("off") ax.grid(False) for p_idx in range(10): ax.scatter( [tensors["player_xs"][time_step, p_idx]], [tensors["player_ys"][time_step, p_idx]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c=colors[p_idx], marker=markers[p_idx], edgecolors="black", ) plt.xlim(auto=False) plt.ylim(auto=False) ax.plot( tensors["player_xs"][:time_step + 1, pred_player_idx], tensors["player_ys"][:time_step + 1, pred_player_idx], c="red", ) ax.plot( tensors["player_xs"][time_step:, pred_player_idx], tensors["player_ys"][time_step:, pred_player_idx], linestyle=":", c="red", ) ax.scatter( [tensors["ball_xs"][time_step]], [tensors["ball_ys"][time_step]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c="#ff8c00", edgecolors="black", ) plt.subplots_adjust(0, 0, 1, 1) fig.canvas.draw() court_img = Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb()) plt.cla() traj_data = np.full((player_traj_n, player_traj_n, 3), 255, dtype=np.uint8) player_traj = tensors["player_trajs"][time_step, pred_player_idx] traj_data[player_traj // player_traj_n, player_traj % player_traj_n] = 0 traj_img = Image.fromarray(traj_data).resize( (scale * player_traj_n, scale * player_traj_n), resample=0) add_grid(traj_img, player_traj_n) traj_pred = np.zeros((player_traj_n, player_traj_n, 3)) traj_pred[:, :] = pred_player_probs[time_step].view( player_traj_n, player_traj_n, 1) pred_img = Image.fromarray(np.uint8(255 - 255 * traj_pred)).resize( (scale * player_traj_n, scale * player_traj_n), resample=0) add_grid(pred_img, player_traj_n) final_img = Image.new( "RGBA", ( court_img.size[0], court_img.size[1] + scale * player_traj_n + vert_buffer, ), ) final_img.paste(court_img, (0, 0)) full_width = 3 * traj_img.size[0] horiz_buffer = court_img.size[0] // 2 - full_width // 2 final_img.paste(traj_img, (horiz_buffer, court_img.size[1] + vert_buffer)) final_img.paste( pred_img, (2 * traj_img.size[0] + horiz_buffer, court_img.size[1] + vert_buffer), ) time_step = str(time_step).zfill(2) final_img.save(f"{home_dir}/test/{test_idx}_{time_step}.png") saved += 1 shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test") shutil.rmtree(f"{home_dir}/test")
def get_nearest_neighbors(): JOB = "20210124114950" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) train_dataset = init_datasets(opts)[0] model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model.eval() player_embeddings = model.player_embedding.weight with torch.no_grad(): if opts["model"]["sigmoid"] == "logistic": player_embeddings = torch.sigmoid(player_embeddings) elif opts["model"]["sigmoid"] == "tanh": player_embeddings = torch.tanh(player_embeddings) player_embeddings = player_embeddings.detach().numpy() all_player_dists = { "cosine": pairwise_distances(player_embeddings, metric="cosine"), "euclidean": pairwise_distances(player_embeddings, metric="euclidean"), } players = [ "LeBron James", "Stephen Curry", "Russell Westbrook", "James Harden", "Pau Gasol", "Anthony Davis", "Serge Ibaka", "Giannis Antetokounmpo", "Kawhi Leonard", "Jimmy Butler", ] k = 20 min_playing_time = 60000 # 0/13314/39917/60000 --> 100%/75%/50%/starters. baller2vec_config = pickle.load( open(f"{DATA_DIR}/baller2vec_config.pydict", "rb")) player_idx2props = baller2vec_config["player_idx2props"] for player in players: print(f"\n{player}") player_idx = get_player_idx(player, player_idx2props) for metric in ["cosine", "euclidean"]: print(f"\n{metric}\n") player_dists = all_player_dists[metric] sim_player_idxs = np.argsort(player_dists[player_idx]) neighbors = 0 for sim_player_idx in sim_player_idxs[1:]: if sim_player_idx in player_idx2props: if "playing_time" not in player_idx2props[sim_player_idx]: continue if (player_idx2props[sim_player_idx]["playing_time"] < min_playing_time): continue sim_player_score = player_dists[player_idx, sim_player_idx] name = player_idx2props.get(sim_player_idx, {"name": "Baller"})["name"] print(f"{name}: {sim_player_score:.4f}") neighbors += 1 if neighbors == k: break keep_embeddings = [] player_names = [] for (player_idx, player_embedding) in enumerate(player_embeddings): if player_idx != len(player_embeddings) - 1: if "playing_time" not in player_idx2props[player_idx]: continue if player_idx2props[player_idx]["playing_time"] < min_playing_time: continue keep_embeddings.append(player_embedding) if player_idx == len(player_embeddings) - 1: player_names.append("Baller") else: player_names.append(player_idx2props[player_idx]["name"]) keep_embeddings = np.vstack(keep_embeddings) home_dir = os.path.expanduser("~") os.makedirs(f"{home_dir}/test", exist_ok=True) np.savetxt(f"{home_dir}/test/embeddings.tsv", keep_embeddings, delimiter="\t") with open(f"{home_dir}/test/names.tsv", "w") as f: print("\n".join(player_names), file=f) shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test") shutil.rmtree(f"{home_dir}/test")
def plot_ball_traj_preds_specific_game(): JOB = "20210124114950" JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}" # Load model. opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml")) (train_dataset, _, _, _, _, _) = init_datasets(opts) model = init_model(opts, train_dataset) model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth")) model.eval() court = plt.imread("court.png") width = 5 height = width * court.shape[0] / float(court.shape[1]) vert_buffer = 10 pt_scale = 1.4 scale = 4 home_dir = os.path.expanduser("~") os.makedirs(f"{home_dir}/test", exist_ok=True) (fig, ax) = plt.subplots(figsize=(width, height)) save_output = SaveOutput() hook_handles = [] for layer in model.transformer.layers: hook_handles.append(layer.self_attn.register_forward_hook(save_output)) gameid = "0021500622" X = np.load(f"{GAMES_DIR}/{gameid}_X.npy") y = np.load(f"{GAMES_DIR}/{gameid}_y.npy") period = 3 period_time = "1:55" start = get_start_idx_from_time(X, period, period_time) tensors = train_dataset.get_sample(X, y, start) sub_rubio = False if sub_rubio: baller2vec_config = pickle.load( open(f"{DATA_DIR}/baller2vec_config.pydict", "rb")) player_idx2props = baller2vec_config["player_idx2props"] lebron_idx = get_player_idx("LeBron James", player_idx2props) rubio_idx = get_player_idx("Ricky Rubio", player_idx2props) player_idxs = tensors["player_idxs"] player_idxs[player_idxs == lebron_idx] = rubio_idx with torch.no_grad(): ball_trajs = model(tensors)["ball"][-len(tensors["player_idxs"]):] ball_probs = torch.softmax(ball_trajs, dim=1) ball_traj_n = train_dataset.ball_traj_n ball_probs = ball_probs.reshape(len(ball_probs), ball_traj_n, ball_traj_n, ball_traj_n) ball_xy_probs = ball_probs.sum(dim=1).permute(0, 2, 1) colors = [] markers = [] for p_idx in range(10): if tensors["player_hoop_sides"][0, p_idx]: colors.append("white") markers.append("s") else: colors.append("gray") markers.append("^") ball_start = 10 * len(tensors["player_idxs"]) for time_step in range(19): print(time_step, flush=True) ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN]) ax.axis("off") ax.grid(False) for p_idx in range(10): ax.scatter( [tensors["player_xs"][time_step, p_idx]], [tensors["player_ys"][time_step, p_idx]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c=colors[p_idx], marker=markers[p_idx], edgecolors="black", ) ax.scatter( [tensors["ball_xs"][time_step]], [tensors["ball_ys"][time_step]], s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2, c="#ff8c00", edgecolors="black", ) plt.subplots_adjust(0, 0, 1, 1) fig.canvas.draw() court_img = Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb()) plt.cla() traj_data = np.full((ball_traj_n, ball_traj_n, 3), 255, dtype=np.uint8) ball_traj = tensors["ball_trajs"][time_step] (ball_row, ball_col_dep) = ( ball_traj // ball_traj_n**2, ball_traj % ball_traj_n**2, ) (ball_col, ball_dep) = (ball_col_dep // ball_traj_n, ball_col_dep % ball_traj_n) traj_data[ball_dep, ball_col] = 0 traj_img = Image.fromarray(traj_data).resize( (scale * ball_traj_n, scale * ball_traj_n), resample=0) add_grid(traj_img, ball_traj_n) traj_pred = np.zeros((ball_traj_n, ball_traj_n, 3)) traj_pred[:, :] = ball_xy_probs[time_step].view( ball_traj_n, ball_traj_n, 1) pred_img = Image.fromarray(np.uint8(255 - 255 * traj_pred)).resize( (scale * ball_traj_n, scale * ball_traj_n), resample=0) add_grid(pred_img, ball_traj_n) final_img = Image.new( "RGBA", ( court_img.size[0], court_img.size[1] + scale * ball_traj_n + vert_buffer, ), ) final_img.paste(court_img, (0, 0)) full_width = 3 * traj_img.size[0] horiz_buffer = court_img.size[0] // 2 - full_width // 2 final_img.paste(traj_img, (horiz_buffer, court_img.size[1] + vert_buffer)) final_img.paste( pred_img, (2 * traj_img.size[0] + horiz_buffer, court_img.size[1] + vert_buffer), ) true_prob = traj_pred[ball_dep, ball_col, 0] print(true_prob) true_prob = f"{true_prob:.2f}".split(".")[1] time_step = str(time_step).zfill(2) final_img.save(f"{home_dir}/test/court_{time_step}_{true_prob}.png") shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test") shutil.rmtree(f"{home_dir}/test")