MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / reset_batch_size

Method reset_batch_size

imperative/python/megengine/utils/network.py:350–376  ·  view source on GitHub ↗

r"""Helper for reset batch size; first dimension of all data providers not in blacklist are assumed to be the batch size Args: blacklist: data provider names whose first dimension is not batchbatch size

(self, batchsize, *, blacklist=())

Source from the content-addressed store, hash-verified

348 i.name = v1
349
350 def reset_batch_size(self, batchsize, *, blacklist=()):
351 r"""Helper for reset batch size; first dimension of all data providers
352 not in blacklist are assumed to be the batch size
353
354 Args:
355 blacklist: data provider names whose first dimension is not
356 batchbatch size
357 """
358 blacklist = set(blacklist)
359 prev_batchsize = None
360 for i in self.data_providers_filter:
361 if i.name in blacklist:
362 blacklist.remove(i.name)
363 else:
364 shp = list(i.shape)
365 if prev_batchsize is None:
366 prev_batchsize = shp[0]
367 else:
368 assert prev_batchsize == shp[0], (
369 "batchsize mismatch: batchsize={} "
370 "shape={} dp={}".format(prev_batchsize, shp, i.name)
371 )
372 shp[0] = batchsize
373 i.shape = tuple(shp)
374 self._compile()
375 assert prev_batchsize is not None, "no data provider found"
376 assert not blacklist, "unused items in blacklist: {}".format(blacklist)
377
378 def replace_vars(self, repl_dict: Dict[VarNode, VarNode]):
379 r"""Replaces vars in the graph.

Callers 2

visualizeFunction · 0.80
test_reset_batchsizeFunction · 0.80

Calls 4

_compileMethod · 0.95
listFunction · 0.85
removeMethod · 0.45
formatMethod · 0.45

Tested by 1

test_reset_batchsizeFunction · 0.64