| 32 | |
| 33 | |
| 34 | def tiled_inference(self, forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype): |
| 35 | # Call y=forward_fn(x) for each tile |
| 36 | tile_num = model_input.shape[-1] |
| 37 | model_output_stack = [] |
| 38 | |
| 39 | for tile_id in range(0, tile_num, tile_batch_size): |
| 40 | |
| 41 | # process input |
| 42 | tile_id_ = min(tile_id + tile_batch_size, tile_num) |
| 43 | x = model_input[:, :, :, :, tile_id: tile_id_] |
| 44 | x = x.to(device=inference_device, dtype=inference_dtype) |
| 45 | x = rearrange(x, "b c h w n -> (n b) c h w") |
| 46 | |
| 47 | # process output |
| 48 | y = forward_fn(x) |
| 49 | y = rearrange(y, "(n b) c h w -> b c h w n", n=tile_id_-tile_id) |
| 50 | y = y.to(device=tile_device, dtype=tile_dtype) |
| 51 | model_output_stack.append(y) |
| 52 | |
| 53 | model_output = torch.concat(model_output_stack, dim=-1) |
| 54 | return model_output |
| 55 | |
| 56 | |
| 57 | def io_scale(self, model_output, tile_size): |