| |
| |
|
|
| import math |
|
|
| import torch |
| from torch.autograd import gradcheck |
|
|
|
|
| from nemo_retriever_ocr.inference.post_processing.rrect_to_quads import RRectToQuads |
|
|
|
|
| def get_rrects(b, h, w): |
| rrects = torch.rand(b, h, w, 5, dtype=torch.float64) |
| rrects[:, :, :, :4] *= 10 |
| rrects[:, :, :, 4] *= 2 * math.pi |
| return rrects |
|
|
|
|
| rrects = get_rrects(2, 5, 5) |
|
|
| cell_size = 4 |
| r2q = RRectToQuads(cell_size) |
|
|
| quads = r2q(rrects) |
|
|
| rrects.requires_grad_() |
|
|
| print("check CPU gradients") |
| gradcheck(r2q.forward, rrects) |
|
|
| rrects = get_rrects(4, 10, 10) |
| rrects.requires_grad_() |
|
|
| print("check GPU gradients") |
| gradcheck(r2q.forward, rrects) |
|
|
