| import webdataset as wds | |
| import os | |
| import torch | |
| class ActivationsDataloader: | |
| def __init__(self, paths_to_datasets, block_name, batch_size, output_or_diff='diff', num_in_buffer=50): | |
| assert output_or_diff in ['diff', 'output'], "Provide 'output' or 'diff'" | |
| self.dataset = wds.WebDataset( | |
| [os.path.join(path_to_dataset, f"{block_name}.tar") | |
| for path_to_dataset in paths_to_datasets] | |
| ).decode("torch") | |
| self.iter = iter(self.dataset) | |
| self.buffer = None | |
| self.pointer = 0 | |
| self.num_in_buffer = num_in_buffer | |
| self.output_or_diff = output_or_diff | |
| self.batch_size = batch_size | |
| self.one_size = None | |
| def renew_buffer(self, to_retrieve): | |
| to_merge = [] | |
| if self.buffer is not None and self.buffer.shape[0] > self.pointer: | |
| to_merge = [self.buffer[self.pointer:].clone()] | |
| del self.buffer | |
| for _ in range(to_retrieve): | |
| sample = next(self.iter) | |
| latents = sample['output.pth'] if self.output_or_diff == 'output' else sample['diff.pth'] | |
| latents = latents.permute((0, 1, 3, 4, 2)) | |
| latents = latents.reshape((-1, latents.shape[-1])) | |
| to_merge.append(latents.to('cuda')) | |
| self.one_size = latents.shape[0] | |
| self.buffer = torch.cat(to_merge, dim=0) | |
| shuffled_indices = torch.randperm(self.buffer.shape[0]) | |
| self.buffer = self.buffer[shuffled_indices] | |
| self.pointer = 0 | |
| def iterate(self): | |
| while True: | |
| if self.buffer == None or self.buffer.shape[0] - self.pointer < self.num_in_buffer * self.one_size * 4 // 5: | |
| try: | |
| to_retrieve = self.num_in_buffer if self.buffer is None else self.num_in_buffer // 5 | |
| self.renew_buffer(to_retrieve) | |
| except StopIteration: | |
| break | |
| batch = self.buffer[self.pointer: self.pointer + self.batch_size] | |
| self.pointer += self.batch_size | |
| assert batch.shape[0] == self.batch_size | |
| yield batch | |