| 215 | |
| 216 | @torch.no_grad() |
| 217 | def init_from_ckpt(self, path, ignore_keys=list(), only_model=False): |
| 218 | sd = torch.load(path, map_location="cpu") |
| 219 | if "state_dict" in list(sd.keys()): |
| 220 | sd = sd["state_dict"] |
| 221 | keys = list(sd.keys()) |
| 222 | for k in keys: |
| 223 | for ik in ignore_keys: |
| 224 | if k.startswith(ik): |
| 225 | print("Deleting key {} from state_dict.".format(k)) |
| 226 | del sd[k] |
| 227 | if self.make_it_fit: |
| 228 | n_params = len([name for name, _ in |
| 229 | itertools.chain(self.named_parameters(), |
| 230 | self.named_buffers())]) |
| 231 | for name, param in tqdm( |
| 232 | itertools.chain(self.named_parameters(), |
| 233 | self.named_buffers()), |
| 234 | desc="Fitting old weights to new weights", |
| 235 | total=n_params |
| 236 | ): |
| 237 | if not name in sd: |
| 238 | continue |
| 239 | old_shape = sd[name].shape |
| 240 | new_shape = param.shape |
| 241 | assert len(old_shape) == len(new_shape) |
| 242 | if len(new_shape) > 2: |
| 243 | # we only modify first two axes |
| 244 | assert new_shape[2:] == old_shape[2:] |
| 245 | # assumes first axis corresponds to output dim |
| 246 | if not new_shape == old_shape: |
| 247 | new_param = param.clone() |
| 248 | old_param = sd[name] |
| 249 | if len(new_shape) == 1: |
| 250 | for i in range(new_param.shape[0]): |
| 251 | new_param[i] = old_param[i % old_shape[0]] |
| 252 | elif len(new_shape) >= 2: |
| 253 | for i in range(new_param.shape[0]): |
| 254 | for j in range(new_param.shape[1]): |
| 255 | new_param[i, j] = old_param[i % old_shape[0], j % old_shape[1]] |
| 256 | |
| 257 | n_used_old = torch.ones(old_shape[1]) |
| 258 | for j in range(new_param.shape[1]): |
| 259 | n_used_old[j % old_shape[1]] += 1 |
| 260 | n_used_new = torch.zeros(new_shape[1]) |
| 261 | for j in range(new_param.shape[1]): |
| 262 | n_used_new[j] = n_used_old[j % old_shape[1]] |
| 263 | |
| 264 | n_used_new = n_used_new[None, :] |
| 265 | while len(n_used_new.shape) < len(new_shape): |
| 266 | n_used_new = n_used_new.unsqueeze(-1) |
| 267 | new_param /= n_used_new |
| 268 | |
| 269 | sd[name] = new_param |
| 270 | |
| 271 | missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( |
| 272 | sd, strict=False) |
| 273 | print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") |
| 274 | if len(missing) > 0: |