utils

• data
  • TensorListDataset
  • dataset_to_tensor()
  • sample_batch()
  • split_dataset()
  • get_class_subset()
• fim
  • BaseKFAC
  • KFAC
  • EKFAC
  • BaseState
  • KFACState
• logger
  • SmoothedValue
  • MetricLogger
  • AverageMeter
• model
  • init_weights()
  • get_layer_name()
  • get_all_layer()
  • get_layer()
  • summary()
  • activate_params()
  • accuracy()
  • generate_target()
  • ExponentialMovingAverage
• module
  • BasicObject
  • Process
  • ModelProcess
  • Module
  • Param
• tensor
  • tanh_func()
  • atan_func()
  • repeat_to_batch()
  • add_noise()
• train
  • train()
  • validate()

class trojanzoo.utils.lock.Lock

A boolean lock class used for contextmanager. It’s used in KFAC to avoid auxiliary computation operations.

Example:

>>> from trojanzoo.utils.lock import Lock
>>>
>>> track = Lock()
>>> print(bool(track))
False
>>> with track():
>>>     print(bool(track))
True
>>> print(bool(track))
False
>>> track.enable()
>>> print(bool(track))
True
>>> track.disable()
>>> print(bool(track))
False

disable()

Set lock boolean value as False. It’s used together with enable() when contextmanager is not suitable for the case.

enable()

Set lock boolean value as True. It’s used together with disable() when contextmanager is not suitable for the case.

trojanzoo.utils.memory.empty_cache(threshold=None)

Call torch.cuda.empty_cache to empty GPU cache when torch.cuda.memory_cached() > thresholdMB.

Parameters:

threshold (float) – The cached memory threshold (MB). Defaults to env['cache_threshold'].

trojanzoo.utils.memory.output_memory(device=None, full=False, indent=0, **kwargs)

Output memory usage information.

Parameters:

• device (None | str | int | device) – Passed to torch.cuda.memory_summary or torch.cuda.memory_allocated | torch.cuda.memory_reserved. Defaults to None.

• full (bool) – Whether to call torch.cuda.memory_summary. Otherwise, call torch.cuda.memory_allocated and torch.cuda.memory_reserved. Defaults to False.

• indent (int) – The space indent for the entire string. Defaults to 0.

• **kwargs – Keyword arguments passed to torch.cuda.memory_summary.