(model, all_params)
| 1273 | |
| 1274 | |
| 1275 | def _RemapParameterBlobsForSharedModel(model, all_params): |
| 1276 | assert model._shared_model |
| 1277 | master_prefix = "{}_{}/".format( |
| 1278 | model._device_prefix, model._devices[0]) |
| 1279 | log.info("Remapping param blobs to master -> {}".format(master_prefix)) |
| 1280 | master_params = set(model.GetParams()) |
| 1281 | |
| 1282 | # Remove all but master params |
| 1283 | def modify_ops(net): |
| 1284 | ops = [] |
| 1285 | for op in net.Proto().op: |
| 1286 | delete_op = False |
| 1287 | # Delete ops that output non-master version of parameter |
| 1288 | for outp in op.output: |
| 1289 | if outp in all_params and outp not in master_params: |
| 1290 | delete_op = True |
| 1291 | log.debug("Delete b/c {}: {}".format(outp, str(op))) |
| 1292 | break |
| 1293 | if delete_op: |
| 1294 | continue |
| 1295 | # Remap inputs to point to the master param |
| 1296 | for j, inp in enumerate(op.input): |
| 1297 | if inp in all_params and inp not in master_params: |
| 1298 | op.input[j] = master_prefix + stripBlobName(inp) |
| 1299 | ops.append(op) |
| 1300 | del net.Proto().op[:] |
| 1301 | net.Proto().op.extend(ops) |
| 1302 | |
| 1303 | modify_ops(model.param_init_net) |
| 1304 | modify_ops(model.net) |
| 1305 | |
| 1306 | |
| 1307 | class CollectivesConcurrencyControl: |
no test coverage detected
searching dependent graphs…