| 80 | return maps |
| 81 | |
| 82 | def create_memory(engine, name, buf, mem, batchsize, inp, inp_idx): |
| 83 | binding_idx = engine.get_binding_index(name) |
| 84 | if binding_idx == -1: |
| 85 | raise AttributeError("Not a valid binding") |
| 86 | print("Binding: name={}, bindingIndex={}".format(name, str(binding_idx))) |
| 87 | dims = engine.get_binding_dimensions(binding_idx).to_DimsCHW() |
| 88 | eltCount = dims.C() * dims.H() * dims.W() * batchsize |
| 89 | |
| 90 | if engine.binding_is_input(binding_idx): |
| 91 | h_mem = inp[inp_idx] |
| 92 | inp_idx = inp_idx + 1 |
| 93 | else: |
| 94 | h_mem = np.random.uniform(0.0, 255.0, eltCount).astype(np.dtype('f4')) |
| 95 | |
| 96 | d_mem = cuda.mem_alloc(eltCount * 4) |
| 97 | cuda.memcpy_htod(d_mem, h_mem) |
| 98 | buf.insert(binding_idx, int(d_mem)) |
| 99 | mem.append(d_mem) |
| 100 | return inp_idx |
| 101 | |
| 102 | |
| 103 | #Run inference on device |