title: How to disable the warning "tensorflow:Method (on_train_batch_end) is slow compared to the batch update (). Check your callbacks"
body: <p>I am trying to implement <a href="https://github.com/timgaripov/swa" rel="nofollow noreferrer">Stochastic Weight Averaging (SWA)</a> with tensorflow 2.0 in a keras style, so I need to update the SWA model weights every step. I have written a custom Callback to do this, but I get the warning every step. Here are some details:</p>
  
  <p>My custom Callback:</p>
  
  <pre class="lang-py prettyprint-override"><code>
  class CustomCallback(tf.keras.callbacks.Callback):
      def __init__(self, valid_data, output_path, swa_alpha=0.99, eval_every=500, eval_batch=16, fold=None):
          self.valid_inputs = valid_data[0]
          self.valid_outputs = valid_data[1]
          self.eval_batch = eval_batch
          self.swa_alpha = swa_alpha
          self.fold = fold
          self.output_path = output_path
          self.rho_value = -1  # record the best rho for report
          self.eval_every = eval_every
  
      def on_train_begin(self, logs={}):
          self.swa_weights = self.model.get_weights()
  
      def on_batch_end(self, batch, logs={}):
  
          # update swa parameters
          alpha = min(1 - 1 / (batch + 1), self.swa_alpha)
          current_weights = self.model.get_weights()
          for i, layer in enumerate(self.model.layers):
              self.swa_weights[i] = alpha * self.swa_weights[i] + (1 - alpha) * current_weights[i]
  
          # validation
          if batch &gt; 0 and batch % self.eval_every == 0:
              # do validation
              val_pred = self.model.predict(self.valid_inputs, batch_size=self.eval_batch)
              rho_val = compute_spearmanr(self.valid_outputs, val_pred)  # the metric
  
              # set the swa parameters and do validation
              self.model.set_weights(self.swa_weights)
              swa_val_pred = self.model.predict(self.valid_inputs, batch_size=self.eval_batch)
              swa_rho_val = compute_spearmanr(self.valid_outputs, swa_val_pred)
  
              # reset the original parameters
              self.model.set_weights(current_weights)
  
              # check whether to save model and update best rho value
              if rho_val &gt; self.rho_value:
                  self.rho_value = rho_val
                  self.model.save_weights(f'{self.output_path}/fold-{fold}-best.h5')
  
          del current_weights
          gc.collect()
  </code></pre>
  
  <p>The output is something like this:</p>
  
  <pre><code>WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.428264). Check your callbacks.
  WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.464315). Check your callbacks.
  WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.502968). Check your callbacks.
  WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.518413). Check your callbacks.
  </code></pre>
  
  <p>I am getting the warning every step, which means without running the validation code, the code for updating SWA params (<code>self.model.get_weights()</code> and the following <code>for</code> loop) is slow enough.</p>
  
  <p>I understand that updating the parameters is very slow because <code>model.get_weights()</code> and <code>model.set_weights()</code> will both do a deep copy of params (a new list of new numpy ndarray according to my experiment).</p>
  
  <p>I think there is nothing wrong with my implementation of SWA (please let me know if there is any bug), so I just want to disable the warning. </p>
  
  <p>What I have tried:</p>
  
  <ol>
  <li>Adding the code <code>os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"</code> and <code>os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"</code> to disable WARNING.</li>
  <li>Setting <code>verbose</code> to <code>2</code> and <code>0</code> in <code>model.fit()</code>, i.e. <code>model.fit(..., verbose=2, ...)</code> and  <code>model.fit(..., verbose=0, ...)</code></li>
  </ol>
  
  <p>Both do no work.</p>
  
  <p>Any idea? Thanks for any help in advance!</p>