| 266 | return torch.from_numpy(np.array(x)).long() |
| 267 | |
| 268 | class EncodeBitMap(torch.nn.Module): |
| 269 | def __init__(self, n=8, fill_value=0.5, ignore_label=0): |
| 270 | super().__init__() |
| 271 | self.n = n |
| 272 | self.fill_value = fill_value |
| 273 | self.ignore_label = ignore_label |
| 274 | |
| 275 | def __call__(self, x): |
| 276 | ignore_mask = x == self.ignore_label |
| 277 | x = torch.bitwise_right_shift(x, torch.arange(self.n, device=x.device)[:, None, None]) |
| 278 | x = torch.remainder(x, 2).float() |
| 279 | x[:, ignore_mask] = self.fill_value |
| 280 | return x, ignore_mask |
| 281 | |
| 282 | def make_train_dataset(args, accelerator): |
| 283 | dataset = load_dataset(f"../dataset/{args.dataset_name}.py", trust_remote_code=True) |
no outgoing calls
no test coverage detected