File size: 7,385 Bytes
23bc32f 5553f7f 23bc32f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
from collections import OrderedDict
import cv2
from pathlib import Path
import random
import shutil
from typing import Callable, Dict
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from src.data import Episode
def configure_optimizer(model: nn.Module, learning_rate: float, weight_decay: float, *blacklist_module_names) -> AdamW:
"""Credits to https://github.com/karpathy/minGPT"""
# separate out all parameters to those that will and won't experience regularizing weight decay
decay = set()
no_decay = set()
whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d)
blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding, nn.Conv2d, nn.GroupNorm)
for mn, m in model.named_modules():
for pn, p in m.named_parameters():
fpn = '%s.%s' % (mn, pn) if mn else pn # full param name
if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
no_decay.add(fpn)
elif 'bias' in pn:
# all biases will not be decayed
no_decay.add(fpn)
elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
# weights of whitelist modules will be weight decayed
decay.add(fpn)
elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
# weights of blacklist modules will NOT be weight decayed
no_decay.add(fpn)
# validate that we considered every parameter
param_dict = {pn: p for pn, p in model.named_parameters()}
inter_params = decay & no_decay
union_params = decay | no_decay
assert len(inter_params) == 0, f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
assert len(param_dict.keys() - union_params) == 0, f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"
# create the pytorch optimizer object
optim_groups = [
{"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": weight_decay},
{"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
]
optimizer = AdamW(optim_groups, lr=learning_rate)
return optimizer
def init_weights(module: nn.Module) -> None:
if isinstance(module, (nn.Linear, nn.Embedding)):
module.weight.data.normal_(mean=0.0, std=0.02)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
def extract_state_dict(state_dict: Dict, module_name: str) -> OrderedDict:
return OrderedDict({k.split('.', 1)[1]: v for k, v in state_dict.items() if k.startswith(module_name)})
def set_seed(seed: int) -> None:
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
random.seed(seed)
@torch.no_grad()
def compute_discounted_returns(rewards: torch.FloatTensor, gamma: float) -> torch.FloatTensor:
assert 0 < gamma <= 1 and rewards.ndim == 2 # (B, T)
gammas = gamma ** torch.arange(rewards.size(1))
r = rewards * gammas
return (r + r.sum(dim=1, keepdim=True) - r.cumsum(dim=1)) / gammas
class LossWithIntermediateLosses:
def __init__(self, **kwargs) -> None:
self.loss_total = sum(kwargs.values())
self.intermediate_losses = {k: v.item() for k, v in kwargs.items()}
class EpisodeDirManager:
def __init__(self, episode_dir: Path, max_num_episodes: int) -> None:
self.episode_dir = episode_dir
self.episode_dir.mkdir(parents=False, exist_ok=True)
self.max_num_episodes = max_num_episodes
self.best_return = float('-inf')
def save(self, episode: Episode, episode_id: int, epoch: int) -> None:
if self.max_num_episodes is not None and self.max_num_episodes > 0:
self._save(episode, episode_id, epoch)
def _save(self, episode: Episode, episode_id: int, epoch: int) -> None:
ep_paths = [p for p in self.episode_dir.iterdir() if p.stem.startswith('episode_')]
assert len(ep_paths) <= self.max_num_episodes
if len(ep_paths) == self.max_num_episodes:
to_remove = min(ep_paths, key=lambda ep_path: int(ep_path.stem.split('_')[1]))
to_remove.unlink()
torch.save(episode.__dict__, self.episode_dir / f'episode_{episode_id}_epoch_{epoch}.pt')
ep_return = episode.compute_metrics().episode_return
if ep_return > self.best_return:
self.best_return = ep_return
path_best_ep = [p for p in self.episode_dir.iterdir() if p.stem.startswith('best_')]
assert len(path_best_ep) in (0, 1)
if len(path_best_ep) == 1:
path_best_ep[0].unlink()
torch.save(episode.__dict__, self.episode_dir / f'best_episode_{episode_id}_epoch_{epoch}.pt')
class RandomHeuristic:
def __init__(self, num_actions):
self.num_actions = num_actions
def act(self, obs):
assert obs.ndim == 4 # (N, H, W, C)
n = obs.size(0)
return torch.randint(low=0, high=self.num_actions, size=(n,))
def make_video(fname, fps, frames):
assert frames.ndim == 4 # (T, H, W, C)
_, h, w, c = frames.shape
assert c == 3
video = cv2.VideoWriter(str(fname), cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
for frame in frames:
video.write(frame[:, :, ::-1])
video.release()
def try_until_no_except(fn: Callable):
while True:
try:
fn()
except:
continue
else:
break
def symlog(x: torch.Tensor) -> torch.Tensor:
return torch.sign(x) * torch.log(torch.abs(x) + 1)
def symexp(x: torch.Tensor) -> torch.Tensor:
return torch.sign(x) * (torch.exp(torch.abs(x)) - 1)
def two_hot(x: torch.FloatTensor, x_min: int = -20, x_max: int = 20, num_buckets: int = 255) -> torch.FloatTensor:
x.clamp_(x_min, x_max - 1e-5)
buckets = torch.linspace(x_min, x_max, num_buckets).to(x.device)
k = torch.searchsorted(buckets, x) - 1
values = torch.stack((buckets[k + 1] - x, x - buckets[k]), dim=-1) / (buckets[k + 1] - buckets[k]).unsqueeze(-1)
two_hots = torch.scatter(x.new_zeros(*x.size(), num_buckets), dim=-1, index=torch.stack((k, k + 1), dim=-1), src=values)
return two_hots
def compute_softmax_over_buckets(logits: torch.FloatTensor, x_min: int = -20, x_max: int = 20, num_buckets: int = 255) -> torch.FloatTensor:
buckets = torch.linspace(x_min, x_max, num_buckets).to(logits.device)
probs = F.softmax(logits, dim=-1)
return probs @ buckets
def plot_counts(counts: np.ndarray) -> Image:
fig, ax = plt.subplots(figsize=(14, 7))
ax.plot(counts)
p = Path('priorities.png')
fig.savefig(p)
plt.close(fig)
im = Image.open(p)
p.unlink()
return im
def compute_mask_after_first_done(ends: torch.LongTensor) -> torch.BoolTensor:
assert ends.ndim == 2
first_one_index = torch.argmax(ends, dim=1)
mask = torch.arange(ends.size(1), device=ends.device).unsqueeze(0) <= first_one_index.unsqueeze(1)
mask = torch.logical_or(mask, ends.sum(dim=1, keepdim=True) == 0)
return mask
|