MCPcopy Create free account
hub / github.com/MotrixLab/ViMoGen / _print_bucket_info

Method _print_bucket_info

datasets/sampler.py:250–301  ·  view source on GitHub ↗
(self, bucket_sample_dict: dict)

Source from the content-addressed store, hash-verified

248 return self.approximate_num_batch
249
250 def _print_bucket_info(self, bucket_sample_dict: dict) -> None:
251 # collect statistics
252 total_samples = 0
253 total_batch = 0
254 num_aspect_dict = defaultdict(lambda: [0, 0])
255 num_hwt_dict = defaultdict(lambda: [0, 0])
256 for k, v in bucket_sample_dict.items():
257 size = len(v)
258 num_batch = size // self.bucket.get_batch_size(k[:-1])
259
260 total_samples += size
261 total_batch += num_batch
262
263 num_aspect_dict[k[-1]][0] += size
264 num_aspect_dict[k[-1]][1] += num_batch
265 num_hwt_dict[k[:-1]][0] += size
266 num_hwt_dict[k[:-1]][1] += num_batch
267
268 # sort
269 num_aspect_dict = dict(
270 sorted(num_aspect_dict.items(), key=lambda x: x[0]))
271 num_hwt_dict = dict(
272 sorted(
273 num_hwt_dict.items(),
274 key=lambda x: (get_num_pixels(x[0][0]), x[0][1]),
275 reverse=True,
276 ))
277 num_hwt_img_dict = {k: v for k, v in num_hwt_dict.items() if k[1] == 1}
278 num_hwt_vid_dict = {k: v for k, v in num_hwt_dict.items() if k[1] > 1}
279
280 # log
281 if dist.get_rank() == 0 and self.verbose:
282 get_logger().info('Bucket Info:')
283 get_logger().info(
284 'Bucket [#sample, #batch] by aspect ratio:\n%s',
285 pformat(num_aspect_dict, sort_dicts=False),
286 )
287 get_logger().info(
288 'Image Bucket [#sample, #batch] by HxWxT:\n%s',
289 pformat(num_hwt_img_dict, sort_dicts=False),
290 )
291 get_logger().info(
292 'Video Bucket [#sample, #batch] by HxWxT:\n%s',
293 pformat(num_hwt_vid_dict, sort_dicts=False),
294 )
295 get_logger().info(
296 '#training batch: %s, #training sample: %s, #non empty bucket: %s',
297 format_numel_str(total_batch),
298 format_numel_str(total_samples),
299 len(bucket_sample_dict),
300 )
301 self.approximate_num_batch = total_batch
302
303 def reset(self):
304 self.last_micro_batch_access_index = 0

Callers 2

__iter__Method · 0.95
get_num_batchMethod · 0.95

Calls 4

get_loggerFunction · 0.90
format_numel_strFunction · 0.90
get_num_pixelsFunction · 0.85
get_batch_sizeMethod · 0.80

Tested by

no test coverage detected