修复gradient accumulation

README.md

@@ -17,7 +17,6 @@
 ## 项目状态
 
 - 目前项目主要架构已经稳定。如发现任何bug或是有功能意见与改进欢迎提交Issue，PR或是联系作者。
-- 如使用梯度积累，loss计算可能存在bug。
 
 ## 使用方法
 
@@ -61,7 +60,7 @@ python ./generate.py --length=50 --nsamples=4 --prefix=xxx --fast_pattern --save
 
 ## FP16与Gradient Accumulation支持
 
-- 我在train.py文件中加入了fp16与gradient accumulation支持，如果你安装了apex并且知道fp16是什么的话，可以修改变量fp16=True来启用。但是目前fp16不收敛，原因不明。
+- 我在train.py文件中加入了fp16与gradient accumulation支持，如果你安装了apex并且知道fp16是什么的话，可以修改变量fp16=True来启用。但是目前fp16可能不收敛，原因不明。
 
 ## 联系作者
 

config/model_config_test.json

@@ -0,0 +1,10 @@
+{
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "n_ctx": 64,
+  "n_embd": 128,
+  "n_head": 2,
+  "n_layer": 1,
+  "n_positions": 64,
+  "vocab_size": 13317
+}

train.py

@@ -207,23 +207,22 @@ def main():
                     torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
 
                 #  optimizer step
-                if (step + 1) % gradient_accumulation == 0:
+                if (overall_step + 1) % gradient_accumulation == 0:
                     running_loss += loss.item()
                     optimizer.step()
                     optimizer.zero_grad()
                     scheduler.step()
-                    overall_step += 1
-                    if (overall_step + 1) % log_step == 0:
-                        tb_writer.add_scalar('loss', loss.item(), overall_step)
                 if (overall_step + 1) % log_step == 0:
+                    tb_writer.add_scalar('loss', loss.item() * gradient_accumulation, overall_step)
                     print('now time: {}:{}. Step {} of piece {} of epoch {}, loss {}'.format(
                         datetime.now().hour,
                         datetime.now().minute,
                         step + 1,
                         piece_num,
                         epoch + 1,
-                        running_loss * gradient_accumulation / log_step))
+                        running_loss * gradient_accumulation / (log_step / gradient_accumulation)))
                     running_loss = 0
+                overall_step += 1
             piece_num += 1
 
         print('saving model for epoch {}'.format(epoch + 1))