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)