Skip to content

Commit

Permalink
Updating docs
Browse files Browse the repository at this point in the history
  • Loading branch information
@samuelstevens
samuelstevens committed Oct 25, 2024
1 parent 974fdf7 commit 3e8132b
Show file tree
Hide file tree
Showing 8 changed files with 260 additions and 383 deletions.
141 changes: 101 additions & 40 deletions docs/saev/config.html
View file

Large diffs are not rendered by default.

60 changes: 38 additions & 22 deletions docs/saev/index.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1">
<meta name="generator" content="pdoc3 0.11.1">
<title>saev API documentation</title>
<meta name="description" content="">
<meta name="description" content="saev is a Python package for training sparse autoencoders (SAEs) on vision transformers (ViTs) in PyTorch …">
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/13.0.0/sanitize.min.css" integrity="sha512-y1dtMcuvtTMJc1yPgEqF0ZjQbhnc/bFhyvIyVNb9Zk5mIGtqVaAB1Ttl28su8AvFMOY0EwRbAe+HCLqj6W7/KA==" crossorigin>
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/13.0.0/typography.min.css" integrity="sha512-Y1DYSb995BAfxobCkKepB1BqJJTPrOp3zPL74AWFugHHmmdcvO+C48WLrUOlhGMc0QG7AE3f7gmvvcrmX2fDoA==" crossorigin>
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/default.min.css" crossorigin>
Expand All @@ -27,6 +27,8 @@
<h1 class="title">Package <code>saev</code></h1>
</header>
<section id="section-intro">
<p>saev is a Python package for training sparse autoencoders (SAEs) on vision transformers (ViTs) in PyTorch.</p>
<p>The main entrypoint to the package is in <a href="https://github.com/samuelstevens/saev/blob/main/main.py">main.py</a>; use <code>python main.py --help</code> to see the options and documentation for the script.</p>
</section>
<section>
<h2 class="section-title" id="header-submodules">Sub-modules</h2>
Expand All @@ -45,20 +47,13 @@ <h2 class="section-title" id="header-submodules">Sub-modules</h2>
</dd>
<dt><code class="name"><a title="saev.modeling" href="modeling.html">saev.modeling</a></code></dt>
<dd>
<div class="desc"></div>
</dd>
<dt><code class="name"><a title="saev.sessions" href="sessions.html">saev.sessions</a></code></dt>
<dd>
<div class="desc"></div>
<div class="desc"><p>modeling is the main module for the saev package and contains all the important non-config classes.
It's fine for this package to be slow to import …</p></div>
</dd>
<dt><code class="name"><a title="saev.training" href="training.html">saev.training</a></code></dt>
<dd>
<div class="desc"></div>
</dd>
<dt><code class="name"><a title="saev.vits" href="vits.html">saev.vits</a></code></dt>
<dd>
<div class="desc"></div>
</dd>
<dt><code class="name"><a title="saev.webapp" href="webapp.html">saev.webapp</a></code></dt>
<dd>
<div class="desc"></div>
Expand All @@ -74,7 +69,7 @@ <h2 class="section-title" id="header-classes">Classes</h2>
<dl>
<dt id="saev.Config"><code class="flex name class">
<span>class <span class="ident">Config</span></span>
<span>(</span><span>image_width: int = 224, image_height: int = 224, model: str = 'ViT-L-14/openai', module_name: str = 'resid', block_layer: int = -2, data: <a title="saev.config.Huggingface" href="config.html#saev.config.Huggingface">Huggingface</a> | <a title="saev.config.Webdataset" href="config.html#saev.config.Webdataset">Webdataset</a> = &lt;factory&gt;, n_workers: int = 8, d_in: int = 1024, n_epochs: int = 3, n_batches_in_store: int = 15, vit_batch_size: int = 1024, expansion_factor: int = 64, l1_coefficient: float = 8e-05, lr: float = 0.0004, lr_warm_up_steps: int = 500, batch_size: int = 1024, use_ghost_grads: bool = True, feature_sampling_window: int = 64, resample_batches: int = 32, feature_reinit_scale: float = 0.2, dead_feature_window: int = 64, dead_feature_estimation_method: str = 'no_fire', dead_feature_threshold: float = 1e-06, log_to_wandb: bool = True, wandb_project: str = 'saev', wandb_log_freq: int = 10, device: str = 'cuda', seed: int = 42, dtype: str = 'float32', checkpoint_path: str = 'checkpoints')</span>
<span>(</span><span>image_width: int = 224, image_height: int = 224, model: str = 'ViT-L-14/openai', module_name: str = 'resid', block_layer: int = -2, data: <a title="saev.config.Imagenet" href="config.html#saev.config.Imagenet">Imagenet</a> | <a title="saev.config.TreeOfLife" href="config.html#saev.config.TreeOfLife">TreeOfLife</a> = &lt;factory&gt;, n_workers: int = 8, d_vit: int = 1024, n_epochs: int = 3, n_batches_in_store: int = 15, vit_batch_size: int = 1024, expansion_factor: int = 64, l1_coefficient: float = 8e-05, lr: float = 0.0004, lr_warm_up_steps: int = 500, batch_size: int = 1024, use_ghost_grads: bool = True, feature_sampling_window: int = 64, resample_batches: int = 32, feature_reinit_scale: float = 0.2, dead_feature_window: int = 64, dead_feature_estimation_method: str = 'no_fire', dead_feature_threshold: float = 1e-06, log_to_wandb: bool = True, wandb_project: str = 'saev', wandb_log_freq: int = 10, device: str = 'cuda', seed: int = 42, dtype: str = 'float32', checkpoint_path: str = 'checkpoints', slurm: bool = False, slurm_acct: str = 'PAS2136', log_to: str = './logs')</span>
</code></dt>
<dd>
<div class="desc"><p>Configuration for training a sparse autoencoder on a vision transformer.</p></div>
Expand All @@ -93,14 +88,16 @@ <h2 class="section-title" id="header-classes">Classes</h2>
image_width: int = 224
image_height: int = 224
model: str = &#34;ViT-L-14/openai&#34;
&#34;&#34;&#34;Model string, for use with open_clip.&#34;&#34;&#34;
module_name: str = &#34;resid&#34;
block_layer: int = -2
data: Huggingface | Webdataset = dataclasses.field(default_factory=Huggingface)
data: Imagenet | TreeOfLife = dataclasses.field(default_factory=Imagenet)
&#34;&#34;&#34;Which dataset to use.&#34;&#34;&#34;
n_workers: int = 8
&#34;&#34;&#34;Number of dataloader workers.&#34;&#34;&#34;

# SAE Parameters
d_in: int = 1024
d_vit: int = 1024

# Activation Store Parameters
n_epochs: int = 3
Expand Down Expand Up @@ -136,13 +133,19 @@ <h2 class="section-title" id="header-classes">Classes</h2>
dtype: str = &#34;float32&#34;
checkpoint_path: str = &#34;checkpoints&#34;

slurm: bool = False
&#34;&#34;&#34;Whether to use submitit to run jobs on a slurm cluster.&#34;&#34;&#34;
slurm_acct: str = &#34;PAS2136&#34;
&#34;&#34;&#34;Slurm account string.&#34;&#34;&#34;
log_to: str = &#34;./logs&#34;

@property
def store_size(self) -&gt; int:
return self.n_batches_in_store * self.batch_size

@property
def d_sae(self) -&gt; int:
return self.d_in * self.expansion_factor
return self.d_vit * self.expansion_factor

@property
def run_name(self) -&gt; str:
Expand All @@ -164,13 +167,13 @@ <h3>Class variables</h3>
<dd>
<div class="desc"></div>
</dd>
<dt id="saev.Config.d_in"><code class="name">var <span class="ident">d_in</span> : int</code></dt>
<dt id="saev.Config.d_vit"><code class="name">var <span class="ident">d_vit</span> : int</code></dt>
<dd>
<div class="desc"></div>
</dd>
<dt id="saev.Config.data"><code class="name">var <span class="ident">data</span> : <a title="saev.config.Huggingface" href="config.html#saev.config.Huggingface">Huggingface</a> | <a title="saev.config.Webdataset" href="config.html#saev.config.Webdataset">Webdataset</a></code></dt>
<dt id="saev.Config.data"><code class="name">var <span class="ident">data</span> : <a title="saev.config.Imagenet" href="config.html#saev.config.Imagenet">Imagenet</a> | <a title="saev.config.TreeOfLife" href="config.html#saev.config.TreeOfLife">TreeOfLife</a></code></dt>
<dd>
<div class="desc"></div>
<div class="desc"><p>Which dataset to use.</p></div>
</dd>
<dt id="saev.Config.dead_feature_estimation_method"><code class="name">var <span class="ident">dead_feature_estimation_method</span> : str</code></dt>
<dd>
Expand Down Expand Up @@ -216,6 +219,10 @@ <h3>Class variables</h3>
<dd>
<div class="desc"></div>
</dd>
<dt id="saev.Config.log_to"><code class="name">var <span class="ident">log_to</span> : str</code></dt>
<dd>
<div class="desc"></div>
</dd>
<dt id="saev.Config.log_to_wandb"><code class="name">var <span class="ident">log_to_wandb</span> : bool</code></dt>
<dd>
<div class="desc"></div>
Expand All @@ -230,7 +237,7 @@ <h3>Class variables</h3>
</dd>
<dt id="saev.Config.model"><code class="name">var <span class="ident">model</span> : str</code></dt>
<dd>
<div class="desc"></div>
<div class="desc"><p>Model string, for use with open_clip.</p></div>
</dd>
<dt id="saev.Config.module_name"><code class="name">var <span class="ident">module_name</span> : str</code></dt>
<dd>
Expand All @@ -256,6 +263,14 @@ <h3>Class variables</h3>
<dd>
<div class="desc"></div>
</dd>
<dt id="saev.Config.slurm"><code class="name">var <span class="ident">slurm</span> : bool</code></dt>
<dd>
<div class="desc"><p>Whether to use submitit to run jobs on a slurm cluster.</p></div>
</dd>
<dt id="saev.Config.slurm_acct"><code class="name">var <span class="ident">slurm_acct</span> : str</code></dt>
<dd>
<div class="desc"><p>Slurm account string.</p></div>
</dd>
<dt id="saev.Config.use_ghost_grads"><code class="name">var <span class="ident">use_ghost_grads</span> : bool</code></dt>
<dd>
<div class="desc"></div>
Expand Down Expand Up @@ -284,7 +299,7 @@ <h3>Instance variables</h3>
</summary>
<pre><code class="python">@property
def d_sae(self) -&gt; int:
return self.d_in * self.expansion_factor</code></pre>
return self.d_vit * self.expansion_factor</code></pre>
</details>
</dd>
<dt id="saev.Config.run_name"><code class="name">prop <span class="ident">run_name</span> : str</code></dt>
Expand Down Expand Up @@ -329,9 +344,7 @@ <h3>Instance variables</h3>
<li><code><a title="saev.config" href="config.html">saev.config</a></code></li>
<li><code><a title="saev.helpers" href="helpers.html">saev.helpers</a></code></li>
<li><code><a title="saev.modeling" href="modeling.html">saev.modeling</a></code></li>
<li><code><a title="saev.sessions" href="sessions.html">saev.sessions</a></code></li>
<li><code><a title="saev.training" href="training.html">saev.training</a></code></li>
<li><code><a title="saev.vits" href="vits.html">saev.vits</a></code></li>
<li><code><a title="saev.webapp" href="webapp.html">saev.webapp</a></code></li>
</ul>
</li>
Expand All @@ -343,8 +356,8 @@ <h4><code><a title="saev.Config" href="#saev.Config">Config</a></code></h4>
<li><code><a title="saev.Config.batch_size" href="#saev.Config.batch_size">batch_size</a></code></li>
<li><code><a title="saev.Config.block_layer" href="#saev.Config.block_layer">block_layer</a></code></li>
<li><code><a title="saev.Config.checkpoint_path" href="#saev.Config.checkpoint_path">checkpoint_path</a></code></li>
<li><code><a title="saev.Config.d_in" href="#saev.Config.d_in">d_in</a></code></li>
<li><code><a title="saev.Config.d_sae" href="#saev.Config.d_sae">d_sae</a></code></li>
<li><code><a title="saev.Config.d_vit" href="#saev.Config.d_vit">d_vit</a></code></li>
<li><code><a title="saev.Config.data" href="#saev.Config.data">data</a></code></li>
<li><code><a title="saev.Config.dead_feature_estimation_method" href="#saev.Config.dead_feature_estimation_method">dead_feature_estimation_method</a></code></li>
<li><code><a title="saev.Config.dead_feature_threshold" href="#saev.Config.dead_feature_threshold">dead_feature_threshold</a></code></li>
Expand All @@ -357,6 +370,7 @@ <h4><code><a title="saev.Config" href="#saev.Config">Config</a></code></h4>
<li><code><a title="saev.Config.image_height" href="#saev.Config.image_height">image_height</a></code></li>
<li><code><a title="saev.Config.image_width" href="#saev.Config.image_width">image_width</a></code></li>
<li><code><a title="saev.Config.l1_coefficient" href="#saev.Config.l1_coefficient">l1_coefficient</a></code></li>
<li><code><a title="saev.Config.log_to" href="#saev.Config.log_to">log_to</a></code></li>
<li><code><a title="saev.Config.log_to_wandb" href="#saev.Config.log_to_wandb">log_to_wandb</a></code></li>
<li><code><a title="saev.Config.lr" href="#saev.Config.lr">lr</a></code></li>
<li><code><a title="saev.Config.lr_warm_up_steps" href="#saev.Config.lr_warm_up_steps">lr_warm_up_steps</a></code></li>
Expand All @@ -368,6 +382,8 @@ <h4><code><a title="saev.Config" href="#saev.Config">Config</a></code></h4>
<li><code><a title="saev.Config.resample_batches" href="#saev.Config.resample_batches">resample_batches</a></code></li>
<li><code><a title="saev.Config.run_name" href="#saev.Config.run_name">run_name</a></code></li>
<li><code><a title="saev.Config.seed" href="#saev.Config.seed">seed</a></code></li>
<li><code><a title="saev.Config.slurm" href="#saev.Config.slurm">slurm</a></code></li>
<li><code><a title="saev.Config.slurm_acct" href="#saev.Config.slurm_acct">slurm_acct</a></code></li>
<li><code><a title="saev.Config.store_size" href="#saev.Config.store_size">store_size</a></code></li>
<li><code><a title="saev.Config.use_ghost_grads" href="#saev.Config.use_ghost_grads">use_ghost_grads</a></code></li>
<li><code><a title="saev.Config.vit_batch_size" href="#saev.Config.vit_batch_size">vit_batch_size</a></code></li>
Expand Down
Loading

0 comments on commit 3e8132b

Please sign in to comment.