def progress(iterable=None, verbose=True, **kwargs): # type: ignore
if not verbose:
return _DummyTqdm(iterable)
Logging.warn(
"`tqdm` package is not installed, no progress bar is shown.",
category=ImportWarning)
if type(iterable) is int:
return range(iterable)
return iterable
def wrapped(model, *args, **kwargs):
# TODO: is there a better way to inspect training loop stats?
nonlocal first_time
try:
result = func(model, *args, **kwargs)
first_time = False
return result
except RuntimeError as e:
if 'out of memory' in str(e):
# this is an OOM error, try releasing memory and do it again
Logging.warn(f"CUDA out of memory error caught. " +
_cuda_memory_str())
gc.collect()
torch.cuda.empty_cache()
Logging.warn(f"Forced GC complete. " + _cuda_memory_str())
# now let's try again
try:
result = func(model, *args, **kwargs)
first_time = False
return result
except RuntimeError as e:
if 'out of memory' in str(e):
if first_time:
# OOM caused by other factors, don't bother saving
raise RuntimeError(
f"CUDA out of memory error caught at first run. "
f"Something else must be wrong.")
else:
# OOM again, this can't be fixed. Save the whole model and optimizer states
filename = f"{model.__class__.__name__}_{time.strftime('%Y-%m-%d %H:%M:%S')}.pt"
save_checkpoint(model, optim, filename=filename)
raise RuntimeError(
f"CUDA out of memory error caught after forced GC. "
f"Model & optimizer saved to {filename}. " +
_cuda_memory_str())
else:
raise e # re-raise the exception
else:
raise e # re-raise the exception
def __init__(self,
num_layers: int,
input_dim: int,
output_dim: int,
hidden_dims: Optional[List[int]] = None,
activation: MaybeList[Union[str, Activation]] = 'id',
bias: MaybeList[bool] = True,
bias_init: Optional[MaybeList[Optional[float]]] = None,
dropout: Optional[MaybeList[float]] = None):
# validate num_layers
if not isinstance(num_layers, int) or num_layers < 1:
raise ValueError("`layers` should be a positive integer.")
# validate hidden_dims
hidden_dims = hidden_dims or []
if len(hidden_dims) != num_layers - 1:
raise ValueError("Length of `hidden_dim` should be `layers` - 1.")
# validate bias
if isinstance(bias, (bool, int)):
bias = [bias] * num_layers
if not (len(bias) == num_layers
and all(isinstance(b, (bool, int)) for b in bias)):
raise ValueError(
"`bias` should be either a boolean, or a list of booleans of length `layers`."
)
# validate bias_init
if bias_init is not None:
if isinstance(bias_init, (float, int)):
bias_init = [bias_init] * num_layers
if not (len(bias_init) == num_layers
and all(b is None or isinstance(b, (float, int))
for b in bias_init)):
raise ValueError(
"`bias_init` should be either a float, or a list of floats of length `layers`."
)
else:
bias_init = [None] * num_layers
# validate dropout
if dropout is not None:
if isinstance(dropout, float):
dropout = [dropout] * num_layers
if not (len(dropout) == num_layers and all(
isinstance(d, float) and 0 <= d < 1 for d in dropout)):
raise ValueError(
"`dropout` should be either a float in range [0, 1),"
" or a list of floats of length `layers`.")
else:
dropout = [0.0] * num_layers
# validate activation
if isinstance(activation, str) or callable(activation):
if activation == 'id' and num_layers > 1:
is_bottleneck = num_layers == 2 and input_dim > hidden_dims[
0] and output_dim > hidden_dims[0]
if not is_bottleneck:
Logging.warn(
"Using identity transform for non-bottleneck MLPs with more than one layer. "
"This is likely an incorrect setting.")
if num_layers == 1:
activation = [activation]
else:
activation = [activation] * (num_layers - 1) + ['id']
elif len(activation) == num_layers - 1:
activation = activation + ['id']
if not (isinstance(activation, list) and len(activation) == num_layers
and all((isinstance(f, str) and f in self.activation_func)
or callable(f) for f in activation)):
raise ValueError(
"Format of `activation` is incorrect. Refer to docstring for details."
)
super().__init__()
dims = [input_dim] + hidden_dims + [output_dim]
self.layers = nn.ModuleList([
nn.Linear(in_features=in_dim, out_features=out_dim, bias=b)
for in_dim, out_dim, b in zip(dims[:-1], dims[1:], bias)
])
self.activations = [
f if callable(f) else self.activation_func[f] for f in activation
]
self.dropouts = [float(d) for d in dropout]
self.bias_init = [
float(b) if b is not None else None for b in bias_init
]
self.reset_parameters()
def wrapped(*args, **kwargs):
Logging.warn(warn_msg, category=DeprecationWarning)
return func(*args, **kwargs)