# Training log of MAE_VSC(Masked Autoencoders with Variational Sparse Coding) 修改與訓練筆記
###### tags: `Training_Log` `ebird` `Masked Autoencoders` `Sparse Coding` `Autoencoder` `Vision Transformer`
:::info
本篇筆記為將Masked Autoencoder的Late
nt vectors(z)加上約束進行稀疏編碼，使z的分布近似為Spike與Slab(多數維度去活化/關閉，使數值為0)，期望模型學習到的表徵/特徵(Representation，例如鳥翅、鳥嘴、鳥羽)之間可以更彼此獨立(可拆解)、集中在少數維度上，更易於解釋，以利於後續生物學上的型態分析
:::
#### 背景知識與提要
- [Masked Autoencoders論文筆記](https://hackmd.io/lTqNcOmQQLiwzkAwVySh8Q)
- [Variational Sparse Coding論文筆記](https://hackmd.io/@YungHuiHsu/HJN5IL2gs)
- [Masked Autoencoders(MAE) 模型訓練筆記](https://hackmd.io/IO3sA_i8QqaM_hScWdscOA)
---
## 一、修改要點
### 1. MAE模型結構修改，加入約束項
見[`mae/models_mae_vsc.py`](https://github.com/YunghuiHsu/ebird_project/blob/main/mae/models_mae_vsc.py)
#### 控制約束強度的主要兩個參數
- alpha($α$, 論文中公式推導使用$γ$)
- Spike and Slab 約束強度。預設0.01
:::spoiler
- 從論文中的測試來看，即使設到相當低的數值(高的約束強度)，模型也會自行修正保留必要維度的資訊，因此在試驗中仍維持0.01這個較高約束強度的設定
:::
- c
- Spike靈敏度的閾值(維度關閉，設為0)，預設50
- 隨著epoch漸增為5至200(as epoch=20,000)
- 在MAE_VSC的epch為1600，訓練結束時，c為50+1600*0.01= 66
- c_delta: c隨著epoch漸增的參數，預設0.01
#### MAE encoder specifics
##### self.norm = norm_layer(embed_dim)
:::spoiler
- 經過這層會將Transformer Block後的z進行對batch之後的維度做標準化
- 剛開始測試時將這層關掉，但發現會導致在訓練很早期就出現loss nan(warmup階段前1/4)因此後來又加回來
- 雖然希望各維度最後達到稀疏分布的效果，但Layer Norm有讓數值正反向傳播時穩定的效果
- 參考資料
- [为什么Transformer要用LayerNorm？](https://www.zhihu.com/question/487766088)
:::
##### self.fc = nn.Linear(embed_dim, 3*embed_dim)
:::spoiler
- 接在Transformer Block模塊後面
- 是可學習的
- inputs : (batch, n_patch*mask_ratio , dim)
- output : (batch, n_patch*mask_ratio , dim * 3)
- 這邊的dim*3是為了後面將z重參數化，控制spike函數活化敏感度
```python
self.fc = nn.Linear(embed_dim, 3*embed_dim)
x = self.fc(x) # (batch, 1+ p^2, embed_dim) > (batch, 1 + p^2, ３*embed_dim)
mu, logvar, logspike = x.chunk(3, dim=-1)
```
:::
#### class MaskedAutoencoderViT(nn.Module)
##### def __init__()
- 初始化設定
:::spoiler
```python=42
def __init__(..., alpha=0.01, c=50, c_delta=0.01):
super().__init__()
# --------------------------------------------------------------------------
# MAE encoder specifics
...
Transformer blocks
self.norm = norm_layer(embed_dim)
self.fc = nn.Linear(embed_dim, 3*embed_dim) # for VSC.
# --------------------------------------------------------------------------
# variational sparse coding(vsc) specifics
# ref: https://paperswithcode.com/paper/variational-sparse-coding
self.alpha = alpha # defaut=0.01
self.c = c # defaut=50
self.c_delta = c_delta # defaut=0.01 for epoch=20,000
```
:::
##### def reparameterize(), def prior_loss()
- `def reparameterize()`
- ELBO RECONSTRUCTION TERM。依照ELB推導的公式將Latent vectors重編碼(Reconstruction term)
- `def prior_loss()`
- ELBO PRIOR TERM。即-(Spike + Slab $D_{KL}$)
- code
:::spoiler
```python=247
# vsc specifics
def reparameterize(self, mu, logvar, logspike):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
gaussian = eps.mul(std).add_(mu)
eta = torch.rand_like(std)
#selection = F.sigmoid(125 * (eta + logspike.exp() - 1))
selection = F.sigmoid(self.c * (eta + logspike.exp() - 1))
return selection.mul(gaussian)
# Reconstruction + KL divergence losses summed over all elements of batch
def prior_loss(self, mu, logvar, logspike):
# see Appendix B from VSC paper / Formula 6
# Calculate each patch, and then calculate the average of the patches and the average of the batches.
spike = torch.clamp(logspike.exp(), 1e-6, 1.0 - 1e-6) # (batch, patch_size_unmasked, embed_dim)
prior1 = -0.5 * torch.sum(spike.mul(1 + logvar - mu.pow(2) - logvar.exp()), dim=-1)
prior21 = (1 - spike).mul(torch.log((1 - spike) / (1 - self.alpha)))
prior22 = spike.mul(torch.log(spike / self.alpha))
prior2 = torch.sum(prior21 + prior22, dim=-1)
prior = prior1 + prior2 # Slab + Spike KL Divergence。 shape : (batch, patch)
loss_prior = prior.mean() # Average loss of the batch
return loss_prior
```
:::
##### def forward_encoder(self, x, mask_ratio)
- 將潛向量(表徵)重參數化使讓其符合想要的分佈reparameterize the Latent vectors (Reconstruction term)
:::spoiler
```python=
def forward_encoder(self, x, mask_ratio):
...
# apply Transformer blocks
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
# vsc specifics ---------------------------------------
# reparameterize for latent vectors
x = self.fc(x) # (batch, 1+ p^2, embed_dim) > (batch, 1 + p^2, ３*embed_dim)
mu, logvar, logspike = x.chunk(3, dim=-1)
logspike = -F.relu(-logspike)
x = self.reparameterize(mu, logvar, logspike)
return x, mask, ids_restore, mu, logvar, logspike
```
:::
##### def forward()
- 在原本forward過程中加入loss_prior，並在函式回傳
:::spoiler
```python=
def forward(self, imgs, mask_ratio=0.75):
# latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
latent, mask, ids_restore , mu, logvar, logspike = self.forward_encoder(imgs, mask_ratio)
loss_prior = self.prior_loss(mu, logvar, logspike)
pred = self.forward_decoder(latent, ids_restore) # [N, L, p*p*3]
loss_rec = self.forward_loss(imgs, pred, mask)
return loss_rec, loss_prior, pred, mask
```
:::
### 2. 訓練腳本與超參數
- [`mae/engine_pretrain_vsc.py`](https://github.com/YunghuiHsu/ebird_project/blob/main/mae/engine_pretrain_vsc.py)
- [`mae/main_pretrain_vsc.py`](https://github.com/YunghuiHsu/ebird_project/blob/main/mae/main_pretrain_vsc.py)
#### 損失函數loss定義與約束強度設定
##### Alpha
- 固定為0.01
##### LR學習率
- 由於在warmup階段經常發生loss_nan，因此將blr從原本 1.5e-4 縮小為1.5e-5後解決部分問題
##### Loss
- 加入loss_prior * 權重
`loss = loss_rec + loss_prior*args.weight_prior`
:::spoiler
- 大原則
- Loss的數值應該以loss_rec為主模型才能有效重建影像
- 過大的loss_prior會導致約束過強，模型無法順利地學習抽象特徵並還原
###### Weight_Prior
- 參考loss_rec 與 loss_prior的級距調整
- 調整至同一級距
- 至少小於0.01
- 使loss_rec > loss_prior
- 試驗測試 至少小於1e-2時不會發生loss_rec NaN
- 數值越小重建影像品質越好
- 1e-3與1e-2時，embedding數值近0的比例均為(84%)
- 1e-7時embedding數值近0的比例為(57%)
:::
##### Loss_Prior Warmup
- 讓loss_prior在warmup階段漸進加入
:::spoiler
- 當loss_prior > loss_rec時，直接縮小loss_prior數值(x0.1)，避免數值過大影響模型參數走向
```python=
if epoch <= args.warmup_epochs:
loss_prior = loss_prior * (epoch / args.warmup_epochs) # 0.0 --> 1.0
# clip loss_prior to less than loss_rec
while loss_rec< loss_prior*args.weight_prior:
loss_prior *= 0.1
```
:::
---
## 二、訓練筆記
### Training log
[Training log on Wandb](https://wandb.ai/yunghui/MAE_VSC_eBird?workspace=user-yunghui)
### 重點摘要
- Loss
- [x] 讓loss_prior < loss_rec
- [x] 在預熱階段將loss_prior漸進加入
- Learning Rate
- [x] 放慢學習速率將預設blr x 0.1倍
- 1.5e-4 -> 1.5e-5
- 配合檢視重建影像，檢視模型學習結果
- 篩選benchmarks影像進行重建，檢視模型學習狀況
:::spoiler
:::
- 其他加速試驗週期迭代的方式
- 抽樣使用小樣本進行模型參數測試
:::spoiler
- 抽樣10%的資料
- 使用更小的抽樣比例會讓魔形無法學習到表徵，無法檢視模型重建還原的圖像
:::
- 其他測試後無效的變量
- Norm_pix_loss、amp off
#### 不同weight_prior與Embedding稀疏性、重建結果
##### 結果
###### 使用eBird前10萬筆資料進行測試
| | w_prior | Sparsity(%)* | lr | epoch | 分布圖 |
|:-------:|:-------:|:------------:|:------:|:-----:|:------------------------------------------:|
| MAE_VSC | 1 | 83.8% | 1.5e-5 | 219^ |  |
| MAE_VSC | 1e-2 | 84.2% | 1.5e-5 | 490 |  |
| MAE_VSC | 1e-3 | 84.4% | 1.5e-5 | 340 |  |
| MAE_VSC | 1e-8 | 57.1% | 1.5e-5 | 300 |  |
| MAE | - | 43.4% | 1.5e-4 | 1174 |  |
- 表格附註
:::spoiler
- ^loss_rec NaN
- *Sparsity : embedding數值近似0的比例( < 1e-7)。使用eBird前10萬筆資料進行測試
- norm_pix_loss=False, amp=False, clip_grad='3'
- 注意y軸比例為log
:::
- 討論
:::spoiler
- weight_prior=1時
- 意即不調整loss_rec與loss_prior比值時，訓練會發生loss_ NaN
- loss_prior比loss_rec 大兩個數量級
- weight_prior在 1e-2 時
- loss_prior與loss_rec約略在同一個數量級，embedding稀疏性與 weight_prior在 1e-3相當
- 與1e-3相比時。loss_rec下降較慢，且重建影像的品質較差
- weight_prior近乎0時(1e-8)
- 仍有稀疏編碼的效果，只是embedding稀疏性明顯較差
- MAE 包括2個分別以0與-15為中心的數值分佈集中處
- 部分維度的數值分佈中心可能並非在0點，造成分析解釋上不易
:::
###### 使用所有資料測試結果(固定間隔取10%)
- w_prior=5e-3, lr=1.5e-5, epoch=400時
- 上圖(logy=True,, xlim=(-30,30) ) 下圖(logy=False, xlim=(-1,1))
### 訓練腳本
:::spoiler
```bash
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 python3 main_pretrain_vsc.py \
--batch_size 400 --accum_iter 10 \
--model mae_vit_base_patch16 \
--mask_ratio 0.75 --blr 1.5e-5 --weight_decay 0.05 \
--epochs 1600 --warmup_epochs 40 \
--data_path /home/esslab/AI_projects/shared/birds \
--alpha 0.01 --weight_prior 5e-3 \
--clip_grad 3 \
--data_file "log_mae_vsc_eBird.txt" \
```
- 如果要用小量資料(10%)測試
- 資料路徑改為讀取隨機抽樣好的csv檔
- `--data_path --data_path birds_sample.csv`
:::
---
## 三、~~各種疑難雜症~~
### 訓練過程loss nan排除
#### 整體問題概述：
- loss = loss_rec + loss_prior* weight
- 在訓練初期及前期經常發生梯度消失loss nan，發生位置在loss_rec
#### 問題推測
- 起始的loss_prior(1e3)與loss_loss(1e-1)不在同一個量級(差距約3-4個量級)
- 過大的loss_prior導致模型在訓練初期或前期無法學習到好的抽象表徵，以至於重建影像(編碼-解碼)過程失敗
- 如果檢視模型重建的影像會得到一片灰階的圖，代表模型在特徵學習與重建任務的參數已壞掉
#### 發生於起始Warmup階段(epoch<5)
:::spoiler
- 問題描述
- warmup初始階段即發生loss_rec nan
- 問題推測
- 學習速率過大
- prior_loss在初期影響太大，導致參數壞掉decoder無法順利重建影像
##### 解決方案1:
- 把 lr 改小 (縮到1/10 )
- warm up 改短 (10個epoch)prior loss 等 warm up 結束後再加入
- [ ] 問題是否解決
- epoch2x時再度發生loss_rec nan
##### 解決方案2: 測試關掉amp(自動混合精度)
- [ ] 問題是否解決
- 在epoch < 5時即發生loss_rec nan
:::
#### 發生於訓練初期(epoch=2x)
:::spoiler
- 問題描述
- 又是loss_rec nan，在lr穩定下降時發生，loss數值看起來都還滿平穩的。
- blr 從1.5e-4 縮小1/10到 1.5e-5
- warmup_epochs =10，在epoch>=10以後，才把loss_prior加入到 loss_total
- 問題推測
- warmup結束後突然把prior加入，是否讓總loss突然跳太快(1e-1 級 跳到1e3，近1000倍)把參數弄壞?
- 這邊的起始loss_priorr就是loss_rec的1000倍
- 解決方案:
-==改成按epoch漸進式把prior_loss加入==，不要讓它一下子跳太快
```python=
if epoch <= args.warmup_epochs:
loss_prior = loss_prior * ((epoch + 1e-7) / args.warmup_epochs) # 0.0 --> 1.0
loss_rec = loss_rec * (10 * (args.warmup_epochs - epoch + 4) / args.warmup_epochs) # 10*1.1 --> 10*0.1
loss = loss_rec + loss_prior
else:
loss = loss_rec + 5*loss_prior
```
- [ ] 問題是否解決
- epoch=15x仍發生loss_rec nan
- 且重建的影像一片模糊沒有物體輪廓
- 上圖是 MAE_VSC版本，epoch=152(發生loss nan前)、下圖是 MAE版本largeepoch=136。
- 下圖MAE版本在epoch136時已經有輪廓了，但MAE_VSC版本卻還是一團糊沒有輪廓
:::
#### 訓練前期(epoch=15x)發生
:::spoiler
- 問題描述
- 問題推測
- loss_prior 數值過大與loss_rec不在同一個數量級，導致prior加入時，讓模型參數過度偏向prior
- 解決策略：
1. 降低vsc的約束強度
- [x] ==調整loss權重 ==
- [ ] 降低編碼稀疏度(提高模型的alpha值)，從0.01改成0.05或0.1
- 解決方案:
- 1. ==在起始階段的數值控制在同一個數量級==，讓warmup階段讓loss_rec漸減、loss_prior漸增，然後放大loss_prior的權重
- 2. 在warmup後，==將loss_prior的權重由5降為2==
```python=
if epoch <= args.warmup_epochs:
loss_prior = loss_prior * ((epoch + 1e-7) / args.warmup_epochs) # 0.0 --> 1.0
loss = loss_rec + loss_prior
else:
loss = loss_rec + 2*loss_prior
```
- [ ] 問題是否解決
- 訓練曲線平順，loss_rec與loss_prior在同一量級，但在warmup階段過後，又得到一片灰白模糊的重建影像
:::
### Decoder重建影像失敗
:::spoiler
#### 問題概述
- 透過warmup階段權重調整，將loss_rec與loss_prior控制在同一量級後，loss nan問題解決，但發現重建的影像變成一團灰階、無法順利重建輪廓
- 上圖：warmup(epoch=40)，結束後，loss_prior變成兩倍後，重建的影像又變成一團糊了(epoch=46)
- 下圖：warmup階段，epoch=26，還有點鳥影
- 在warmup階段透過將loss_rec漸增與loss_prior漸減，讓兩者過渡到同一量級，避免參數被突然過大的loss_prior影響
#### 問題推測
- 稀疏編碼的約束強度過高
#### 解決方案
- 降低vsc的約束強度
- [x] 調整loss權重
- [x] 降低編碼稀疏度(提高模型的alpha值)，從0.01改成 0.1
#### 結果
- warmup結束時(epoch=9)
- 顏色部分不見，但鳥體的部分深色、形體仍保有輪廓
- epoch=12時，loss_rec突然往下掉，重建影像又變一片糊
:::
Sign in with Wallet
Connect another wallet