Move zero_grad

code_samples/base/accelerate

@@ -7,11 +7,11 @@ train_dataloader, model, optimizer scheduler = accelerator.prepare(
 
 model.train()
 for batch in train_dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
     outputs = model(inputs)
     loss = loss_function(outputs, targets)
     accelerator.backward(loss)
     optimizer.step()
     scheduler.step()
+    optimizer.zero_grad()
 </pre>

code_samples/base/basic

@@ -7,7 +7,6 @@
 +)
 
 for batch in dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
 -    inputs = inputs.to(device)
 -    targets = targets.to(device)
@@ -16,7 +15,8 @@ for batch in dataloader:
 -    loss.backward()
 +    accelerator.backward(loss)
     optimizer.step()
-    scheduler.step()</pre>
+    scheduler.step()
+    optimizer.zero_grad()</pre>
 ##
 Everything around `accelerate` occurs with the `Accelerator` class. To use it, first make an object.
 Then call `.prepare` passing in the PyTorch objects that you would normally train with. This will

code_samples/base/calculating_metrics

@@ -11,7 +11,6 @@ import evaluate
 +)
 metric = evaluate.load("accuracy")
 for batch in train_dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
 -    inputs = inputs.to(device)
 -    targets = targets.to(device)
@@ -20,6 +19,7 @@ for batch in train_dataloader:
     loss.backward()
     optimizer.step()
     scheduler.step()
+    optimizer.zero_grad()
 
 model.eval()
 for batch in eval_dataloader:

code_samples/base/checkpointing

@@ -7,13 +7,13 @@ dataloader, model, optimizer scheduler = accelerator.prepare(
 )
 
 for batch in dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
     outputs = model(inputs)
     loss = loss_function(outputs, targets)
     accelerator.backward(loss)
     optimizer.step()
     scheduler.step()
+    optimizer.zero_grad()
 +accelerator.save_state("checkpoint_dir")
 +accelerator.load_state("checkpoint_dir")</pre>
 ##

code_samples/base/experiment_tracking

@@ -9,7 +9,6 @@ train_dataloader, model, optimizer scheduler = accelerator.prepare(
 +accelerator.init_trackers()
 model.train()
 for batch in train_dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
     outputs = model(inputs)
     loss = loss_function(outputs, targets)
@@ -17,6 +16,7 @@ for batch in train_dataloader:
     accelerator.backward(loss)
     optimizer.step()
     scheduler.step()
+    optimizer.zero_grad()
 +accelerator.end_training()
 </pre>
 ##

code_samples/base/gradient_accumulation

@@ -10,13 +10,13 @@ dataloader, model, optimizer scheduler = accelerator.prepare(
 
 for batch in dataloader:
 +  with accelerator.accumulate(model):
-      optimizer.zero_grad()
       inputs, targets = batch
       outputs = model(inputs)
       loss = loss_function(outputs, targets)
       accelerator.backward(loss)
       optimizer.step()
-      scheduler.step()</pre>
+      scheduler.step()
+      optimizer.zero_grad()</pre>
 
 ##
 When performing gradient accumulation in a distributed setup, there are many opportunities for efficiency mistakes

code_samples/base/initial

@@ -1,6 +1,5 @@
 <pre>
 for batch in dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
     inputs = inputs.to(device)
     targets = targets.to(device)
@@ -8,4 +7,5 @@ for batch in dataloader:
     loss = loss_function(outputs, targets)
     loss.backward()
     optimizer.step()
-    scheduler.step()</pre>
+    scheduler.step()
+    optimizer.zero_grad()</pre>

code_samples/base/initial_with_metrics

@@ -2,7 +2,6 @@
 import evaluate
 metric = evaluate.load("accuracy")
 for batch in train_dataloader:
-    optimizer.zero_grad()
     inputs, targets = batch
     inputs = inputs.to(device)
     targets = targets.to(device)
@@ -11,6 +10,7 @@ for batch in train_dataloader:
     loss.backward()
     optimizer.step()
     scheduler.step()
+    optimizer.zero_grad()
 
 model.eval()
 for batch in eval_dataloader: