| | import numpy as np |
| | import pdb |
| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| | |
| |
|
| | class OnlineBuffer(nn.Module): |
| | def __init__(self, buffer_size, batch_size, input_size): |
| | super().__init__() |
| |
|
| | self.place_left = True |
| | self.strategy = None |
| | self.buffer_size = buffer_size |
| | print('buffer has %d slots' % buffer_size, buffer_size) |
| |
|
| | buf_data = torch.FloatTensor(buffer_size, *input_size).fill_(0) |
| | buf_targets = torch.LongTensor(buffer_size).fill_(0) |
| | buf_tasks = torch.LongTensor(buffer_size).fill_(0) |
| |
|
| | self.current_index = 0 |
| | self.n_seen_so_far = 0 |
| | self.is_full = 0 |
| | self.total_classes = 0 |
| | |
| | self.register_buffer('buf_data', buf_data) |
| | self.register_buffer('buf_targets', buf_targets) |
| | self.register_buffer('buf_tasks', buf_tasks) |
| |
|
| |
|
| | def tensor_to_device(self, device): |
| | self.device = device |
| | self.buf_data.to(device) |
| | self.buf_targets.to(device) |
| | self.buf_tasks.to(device) |
| |
|
| |
|
| |
|
| | def add_reservoir(self, x, y, task): |
| | n_elem = x.size(0) |
| | |
| | self.device = x.device |
| | place_left = max(0, self.buffer_size - self.current_index) |
| | offset = min(place_left, n_elem) |
| |
|
| | if place_left: |
| | offset = min(place_left, n_elem) |
| |
|
| | self.buf_data[self.current_index: self.current_index + offset].data.copy_(x[:offset]) |
| | self.buf_targets[self.current_index: self.current_index + offset].data.copy_(y[:offset]) |
| | self.buf_tasks[self.current_index: self.current_index + offset].fill_(task) |
| | self.current_index += offset |
| | self.n_seen_so_far += offset |
| |
|
| | if offset == x.size(0): |
| | return |
| |
|
| | self.place_left = False |
| |
|
| | |
| | x, y = x[place_left:], y[place_left:] |
| |
|
| | indices = torch.FloatTensor(x.size(0)).to(x.device).uniform_(0, self.n_seen_so_far).long() |
| | valid_indices = (indices < self.buf_data.size(0)).long() |
| |
|
| | idx_new_data = valid_indices.nonzero().squeeze(-1) |
| | idx_buffer = indices[idx_new_data] |
| |
|
| | self.n_seen_so_far += x.size(0) |
| |
|
| | if idx_buffer.numel() == 0: |
| | return |
| |
|
| | assert idx_buffer.max() < self.buf_data.size(0), pdb.set_trace() |
| | assert idx_buffer.max() < self.buf_targets.size(0), pdb.set_trace() |
| | assert idx_buffer.max() < self.buf_tasks.size(0), pdb.set_trace() |
| |
|
| | assert idx_new_data.max() < x.size(0), pdb.set_trace() |
| | assert idx_new_data.max() < y.size(0), pdb.set_trace() |
| |
|
| | if self.buf_data.device != x.device: |
| | self.buf_data = self.buf_data.to(x.device) |
| | self.buf_targets = self.buf_targets.to(x.device) |
| | self.buf_tasks = self.buf_tasks.to(x.device) |
| |
|
| | self.buf_data[idx_buffer] = x[idx_new_data] |
| | self.buf_targets[idx_buffer] = y[idx_new_data] |
| | self.buf_tasks[idx_buffer] = task |
| |
|
| |
|
| |
|
| |
|
| | def sample(self, amount, exclude_task = None, ret_ind = False): |
| |
|
| | if self.buf_data.device != self.device: |
| | self.buf_data = self.buf_data.to(self.device) |
| | self.buf_targets = self.buf_targets.to(self.device) |
| | self.buf_tasks = self.buf_tasks.to(self.device) |
| |
|
| | if exclude_task is not None: |
| | valid_indices = (self.t != exclude_task) |
| | valid_indices = valid_indices.nonzero().squeeze() |
| | bx, by, bt = self.buf_data[valid_indices], self.buf_targets[valid_indices], self.buf_tasks[valid_indices] |
| | else: |
| | bx, by, bt = self.buf_data[:self.current_index], self.buf_targets[:self.current_index], self.buf_tasks[:self.current_index] |
| |
|
| | if bx.size(0) < amount: |
| | if ret_ind: |
| | return bx, by, bt, torch.from_numpy(np.arange(bx.size(0))) |
| | else: |
| | return bx, by, bt |
| | else: |
| | indices = torch.from_numpy(np.random.choice(bx.size(0), amount, replace=False)) |
| | indices = indices.to(self.device) |
| |
|
| | if ret_ind: |
| | return bx[indices], by[indices], bt[indices], indices |
| | else: |
| | return bx[indices], by[indices], bt[indices] |
