সাহায্য Kaggle উপর TensorFlow সঙ্গে গ্রেট বেরিয়ার রিফ রক্ষা চ্যালেঞ্জ যোগদান

উন্নত স্বয়ংক্রিয় পার্থক্য

TensorFlow.org এ দেখুন GitHub-এ উৎস দেখুন নোটবুক ডাউনলোড করুন

ভূমিকা গ্রেডিয়েন্ট এবং স্বয়ংক্রিয় বিভেদ করার নির্দেশিকা TensorFlow মধ্যে ক্যালকুলেট গ্রেডিয়েন্ট করা প্রয়োজন সবকিছু অন্তর্ভুক্ত করা হয়েছে। এই নির্দেশিকাটি গভীর, কম সাধারণ বৈশিষ্ট্য উপর গুরুত্ত্ব দেয় tf.GradientTape API- টি।

সেটআপ

import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams['figure.figsize'] = (8, 6)

গ্রেডিয়েন্ট রেকর্ডিং নিয়ন্ত্রণ করা

ইন স্বয়ংক্রিয় বিভেদ নির্দেশিকা কেমন যা ভেরিয়েবল এবং tensors টেপ দ্বারা দেখা হয় গ্রেডিয়েন্ট হিসাব বিল্ডিং যখন নিয়ন্ত্রণ করতে দেখেছি।

টেপে রেকর্ডিং ম্যানিপুলেট করার পদ্ধতিও রয়েছে।

রেকর্ডিং বন্ধ করুন

আপনি গ্রেডিয়েন্ট রেকর্ডিং বন্ধ করতে চান, আপনি ব্যবহার করতে পারেন tf.GradientTape.stop_recording সাময়িকভাবে রেকর্ডিং স্থগিত করার।

আপনি যদি আপনার মডেলের মাঝখানে একটি জটিল অপারেশনকে আলাদা করতে না চান তবে ওভারহেড কমাতে এটি কার্যকর হতে পারে। এটি একটি মেট্রিক বা একটি মধ্যবর্তী ফলাফল গণনা অন্তর্ভুক্ত করতে পারে:

x = tf.Variable(2.0)
y = tf.Variable(3.0)

with tf.GradientTape() as t:
  x_sq = x * x
  with t.stop_recording():
    y_sq = y * y
  z = x_sq + y_sq

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

স্ক্র্যাচ থেকে রেকর্ডিং রিসেট/শুরু করুন

আপনি সম্পূর্ণরূপে নতুন করে শুরু করতে চান, ব্যবহার tf.GradientTape.reset । কেবলমাত্র সাধারণত পড়ার জন্য সহজ গ্রেডিয়েন্ট টেপ ব্লক প্রস্থান এবং পুনরায় চালু করা হয়, কিন্তু আপনি ব্যবহার করতে পারেন reset যখন থেকে প্রস্থান টেপ ব্লক কঠিন বা অসম্ভব পদ্ধতি।

x = tf.Variable(2.0)
y = tf.Variable(3.0)
reset = True

with tf.GradientTape() as t:
  y_sq = y * y
  if reset:
    # Throw out all the tape recorded so far.
    t.reset()
  z = x * x + y_sq

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

নির্ভুলতার সাথে গ্রেডিয়েন্ট প্রবাহ বন্ধ করুন

উপরের গ্লোবাল টেপ নিয়ন্ত্রণের বিপরীতে, tf.stop_gradient ফাংশন অনেক বেশি সুনির্দিষ্ট। এটি একটি নির্দিষ্ট পথ বরাবর প্রবাহিত থেকে গ্রেডিয়েন্টগুলি বন্ধ করতে ব্যবহার করা যেতে পারে, টেপটিতে অ্যাক্সেসের প্রয়োজন ছাড়াই:

x = tf.Variable(2.0)
y = tf.Variable(3.0)

with tf.GradientTape() as t:
  y_sq = y**2
  z = x**2 + tf.stop_gradient(y_sq)

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

কাস্টম গ্রেডিয়েন্ট

কিছু ক্ষেত্রে, আপনি ডিফল্ট ব্যবহার করার পরিবর্তে গ্রেডিয়েন্টগুলি ঠিক কীভাবে গণনা করা হয় তা নিয়ন্ত্রণ করতে চাইতে পারেন। এই পরিস্থিতিতে অন্তর্ভুক্ত:

  1. আপনি লিখছেন একটি নতুন অপের জন্য কোন সংজ্ঞায়িত গ্রেডিয়েন্ট নেই.
  2. ডিফল্ট গণনা সংখ্যাগতভাবে অস্থির।
  3. আপনি ফরোয়ার্ড পাস থেকে একটি ব্যয়বহুল গণনা ক্যাশে করতে চান।
  4. আপনি একটি মান (উদাহরণ স্বরূপ ব্যবহার পরিবর্তন করতে চান tf.clip_by_value বা tf.math.round গ্রেডিয়েন্ট পরিবর্তন ছাড়াই)।

প্রথম ক্ষেত্রে, একটি নতুন অপ আপনি ব্যবহার করতে পারেন লিখতে tf.RegisterGradient আপনার নিজের সেট আপ করার জন্য (বিস্তারিত জন্য API ডক্স দেখুন)। (দ্রষ্টব্য যে গ্রেডিয়েন্ট রেজিস্ট্রি বিশ্বব্যাপী, তাই সাবধানতার সাথে এটি পরিবর্তন করুন।)

আধুনিক তিনটি ক্ষেত্রে, যদি আপনি ব্যবহার করতে পারেন tf.custom_gradient

এখানে প্রযোজ্য একটি উদাহরণ tf.clip_by_norm অন্তর্বর্তী গ্রেডিয়েন্ট করুন:

# Establish an identity operation, but clip during the gradient pass.
@tf.custom_gradient
def clip_gradients(y):
  def backward(dy):
    return tf.clip_by_norm(dy, 0.5)
  return y, backward

v = tf.Variable(2.0)
with tf.GradientTape() as t:
  output = clip_gradients(v * v)
print(t.gradient(output, v))  # calls "backward", which clips 4 to 2
tf.Tensor(2.0, shape=(), dtype=float32)

পড়ুন tf.custom_gradient আরো বিস্তারিত জানার জন্য প্রসাধক এপিআই ডক্স।

SavedModel-এ কাস্টম গ্রেডিয়েন্ট

কাস্টম গ্রেডিয়েন্ট বিকল্প ব্যবহার করে SavedModel করার সংরক্ষণ করা যাবে tf.saved_model.SaveOptions(experimental_custom_gradients=True)

SavedModel মধ্যে সংরক্ষিত করা, গ্রেডিয়েন্ট ফাংশন (আরও জানতে খুঁজে বার করো অনুসরণযোগ্য হওয়া আবশ্যক tf.function সঙ্গে উন্নত কর্মক্ষমতা গাইড)।

class MyModule(tf.Module):

  @tf.function(input_signature=[tf.TensorSpec(None)])
  def call_custom_grad(self, x):
    return clip_gradients(x)

model = MyModule()
tf.saved_model.save(
    model,
    'saved_model',
    options=tf.saved_model.SaveOptions(experimental_custom_gradients=True))

# The loaded gradients will be the same as the above example.
v = tf.Variable(2.0)
loaded = tf.saved_model.load('saved_model')
with tf.GradientTape() as t:
  output = loaded.call_custom_grad(v * v)
print(t.gradient(output, v))
INFO:tensorflow:Assets written to: saved_model/assets
tf.Tensor(2.0, shape=(), dtype=float32)

উপরোক্ত উদাহরণে সম্পর্কে একটি নোট: আপনার সাথে উপরের কোড প্রতিস্থাপন চেষ্টা করুন tf.saved_model.SaveOptions(experimental_custom_gradients=False) , গ্রেডিয়েন্ট এখনও লোড হচ্ছে উপর একই ফলাফল উত্পাদন করা হবে। কারণ গ্রেডিয়েন্ট রেজিস্ট্রি এখনও ফাংশন ব্যবহার করা কাস্টম গ্রেডিয়েন্ট রয়েছে call_custom_op । যাইহোক, যদি আপনি কাস্টম গ্রেডিয়েন্ট ছাড়া সংরক্ষণ অধীন লোড মডেল চলমান পরে রানটাইম পুনর্সূচনা tf.GradientTape ত্রুটি নিক্ষেপ করা হবে: LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN)

একাধিক টেপ

একাধিক টেপ নির্বিঘ্নে যোগাযোগ করে।

উদাহরণস্বরূপ, এখানে প্রতিটি টেপ বিভিন্ন টেনসরের সেট দেখে:

x0 = tf.constant(0.0)
x1 = tf.constant(0.0)

with tf.GradientTape() as tape0, tf.GradientTape() as tape1:
  tape0.watch(x0)
  tape1.watch(x1)

  y0 = tf.math.sin(x0)
  y1 = tf.nn.sigmoid(x1)

  y = y0 + y1

  ys = tf.reduce_sum(y)
tape0.gradient(ys, x0).numpy()   # cos(x) => 1.0
1.0
tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25
0.25

উচ্চ-ক্রম গ্রেডিয়েন্ট

ভেতরে অপারেশনস tf.GradientTape প্রসঙ্গ ম্যানেজার স্বয়ংক্রিয় বিভেদ জন্য রেকর্ড করা হয়। যদি গ্রেডিয়েন্টগুলি সেই প্রসঙ্গে গণনা করা হয়, তবে গ্রেডিয়েন্ট গণনাটিও রেকর্ড করা হয়। ফলস্বরূপ, সঠিক একই API উচ্চ-ক্রম গ্রেডিয়েন্টের জন্যও কাজ করে।

উদাহরণ স্বরূপ:

x = tf.Variable(1.0)  # Create a Tensorflow variable initialized to 1.0

with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    y = x * x * x

  # Compute the gradient inside the outer `t2` context manager
  # which means the gradient computation is differentiable as well.
  dy_dx = t1.gradient(y, x)
d2y_dx2 = t2.gradient(dy_dx, x)

print('dy_dx:', dy_dx.numpy())  # 3 * x**2 => 3.0
print('d2y_dx2:', d2y_dx2.numpy())  # 6 * x => 6.0
dy_dx: 3.0
d2y_dx2: 6.0

যদিও যে আপনি একটি স্কেলার ফাংশন দ্বিতীয় ব্যুৎপন্ন দিতে হয়, তাহলে এই প্যাটার্ন, একটি চট ম্যাট্রিক্স উত্পাদন করতে যেহেতু সাধারণের নেই tf.GradientTape.gradient শুধুমাত্র একটি স্কেলার নতিমাত্রা নির্ণয় করে। একটি গঠন করা চট ম্যাট্রিক্স , এখানে যান চট উদাহরণ অধীনে Jacobian অধ্যায়

"নেস্টেড কল tf.GradientTape.gradient " একটি ভাল প্যাটার্ন যখন আপনি একটি গ্রেডিয়েন্ট থেকে একটি স্কেলার গণক করা হয়, এবং তারপর ফলে স্কালে একটি দ্বিতীয় গ্রেডিয়েন্ট হিসাব জন্য একটি উৎস হিসাবে কাজ করে যা নিম্নলিখিত উদাহরণে হয়।

উদাহরণ: ইনপুট গ্রেডিয়েন্ট নিয়মিতকরণ

অনেক মডেল "বিরোধী উদাহরণ" এর জন্য সংবেদনশীল। কৌশলগুলির এই সংগ্রহটি মডেলের আউটপুটকে বিভ্রান্ত করতে মডেলের ইনপুট পরিবর্তন করে। যেমন সহজ বাস্তবায়ন-যেমন adversarial উদাহরণ ফাস্ট গ্রেডিয়েন্ট পদ্ধতি আক্রমণ ব্যবহার করে স্বাক্ষর করা ইনপুট থেকে সম্মান সঙ্গে আউটপুট গ্রেডিয়েন্ট বরাবর একটি একক পদক্ষেপ -takes; "ইনপুট গ্রেডিয়েন্ট"।

এক কৌশল adversarial উদাহরণ থেকে বলিষ্ঠতার বৃদ্ধি হয় ইনপুট গ্রেডিয়েন্ট নিয়মিতকরণ (ফিনলে & Oberman, 2019), যা প্রচেষ্টা ইনপুট গ্রেডিয়েন্ট মাত্রার কমানোর জন্য। যদি ইনপুট গ্রেডিয়েন্ট ছোট হয়, তাহলে আউটপুটে পরিবর্তনটিও ছোট হওয়া উচিত।

নীচে ইনপুট গ্রেডিয়েন্ট রেগুলারাইজেশনের একটি নির্বোধ বাস্তবায়ন। বাস্তবায়ন হল:

  1. একটি অভ্যন্তরীণ টেপ ব্যবহার করে ইনপুটের সাপেক্ষে আউটপুটের গ্রেডিয়েন্ট গণনা করুন।
  2. সেই ইনপুট গ্রেডিয়েন্টের মাত্রা গণনা করুন।
  3. মডেলের সাপেক্ষে সেই মাত্রার গ্রেডিয়েন্ট গণনা করুন।
x = tf.random.normal([7, 5])

layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)
with tf.GradientTape() as t2:
  # The inner tape only takes the gradient with respect to the input,
  # not the variables.
  with tf.GradientTape(watch_accessed_variables=False) as t1:
    t1.watch(x)
    y = layer(x)
    out = tf.reduce_sum(layer(x)**2)
  # 1. Calculate the input gradient.
  g1 = t1.gradient(out, x)
  # 2. Calculate the magnitude of the input gradient.
  g1_mag = tf.norm(g1)

# 3. Calculate the gradient of the magnitude with respect to the model.
dg1_mag = t2.gradient(g1_mag, layer.trainable_variables)
[var.shape for var in dg1_mag]
[TensorShape([5, 10]), TensorShape([10])]

জ্যাকোবিয়ানরা

পূর্ববর্তী সমস্ত উদাহরণগুলি কিছু উত্স টেনসর(গুলি) এর সাপেক্ষে একটি স্কেলার লক্ষ্যের গ্রেডিয়েন্ট গ্রহণ করেছে।

Jacobian ম্যাট্রিক্স একটি ভেক্টর মূল্যবান ফাংশনের গ্রেডিয়েন্ট প্রতিনিধিত্ব করে। প্রতিটি সারিতে ভেক্টরের উপাদানগুলির একটির গ্রেডিয়েন্ট রয়েছে।

tf.GradientTape.jacobian পদ্ধতি দক্ষতার একটি Jacobian ম্যাট্রিক্স নিরূপণ করতে পারবেন।

মনে রাখবেন যে:

  • ভালো লেগেছে gradient : sources যুক্তি টেন্সর বা tensors একটি ধারক হতে পারে না।
  • ভিন্ন gradient : target টেন্সর একটি একক টেন্সর হতে হবে।

স্কেলার উত্স

প্রথম উদাহরণ হিসাবে, এখানে একটি স্কেলার-উৎস সম্পর্কিত ভেক্টর-টার্গেটের জ্যাকোবিয়ান।

x = tf.linspace(-10.0, 10.0, 200+1)
delta = tf.Variable(0.0)

with tf.GradientTape() as tape:
  y = tf.nn.sigmoid(x+delta)

dy_dx = tape.jacobian(y, delta)

ফলাফলের যখন আপনি একটি স্কেলার থেকে সম্মান সঙ্গে Jacobian নিতে লক্ষ্য আকৃতি আছে, এবং উৎস থেকে সম্মান সঙ্গে প্রতিটি উপাদান গ্রেডিয়েন্ট দেয়:

print(y.shape)
print(dy_dx.shape)
(201,)
(201,)
plt.plot(x.numpy(), y, label='y')
plt.plot(x.numpy(), dy_dx, label='dy/dx')
plt.legend()
_ = plt.xlabel('x')

png

টেনসর উৎস

কিনা ইনপুট স্কালে বা টেন্সর হয়, tf.GradientTape.jacobian দক্ষতার টার্গেট (গুলি) প্রতিটি উপাদান থেকে সম্মান সঙ্গে উৎস প্রতিটি উপাদান গ্রেডিয়েন্ট হিসাব করে।

উদাহরণস্বরূপ, এটি স্তর আউটপুট একটি আকৃতি হয়েছে (10, 7) :

x = tf.random.normal([7, 5])
layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)

with tf.GradientTape(persistent=True) as tape:
  y = layer(x)

y.shape
TensorShape([7, 10])

আর লেয়ারটির কার্নেলের আকৃতি হয় (5, 10) :

layer.kernel.shape
TensorShape([5, 10])

কার্নেলের সাপেক্ষে আউটপুটের জ্যাকোবিয়ানের আকৃতি হল সেই দুটি আকৃতি একত্রে সংযুক্ত:

j = tape.jacobian(y, layer.kernel)
j.shape
TensorShape([7, 10, 5, 10])

আপনি লক্ষ্য এর মাত্রা বেশি যোগফল থাকে, তাহলে আপনি সমষ্টি যে নির্ণিত করা হত নতিমাত্রা অবশিষ্ট করছি tf.GradientTape.gradient :

g = tape.gradient(y, layer.kernel)
print('g.shape:', g.shape)

j_sum = tf.reduce_sum(j, axis=[0, 1])
delta = tf.reduce_max(abs(g - j_sum)).numpy()
assert delta < 1e-3
print('delta:', delta)
g.shape: (5, 10)
delta: 2.3841858e-07

উদাহরণ: হেসিয়ান

যদিও tf.GradientTape একটি নির্মাণের জন্য একটি সুনির্দিষ্ট পদ্ধতি দেয় না চট ম্যাট্রিক্স এটি ব্যবহার এক নির্মাণ করতে সম্ভব tf.GradientTape.jacobian পদ্ধতি।

x = tf.random.normal([7, 5])
layer1 = tf.keras.layers.Dense(8, activation=tf.nn.relu)
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.relu)

with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    x = layer1(x)
    x = layer2(x)
    loss = tf.reduce_mean(x**2)

  g = t1.gradient(loss, layer1.kernel)

h = t2.jacobian(g, layer1.kernel)
print(f'layer.kernel.shape: {layer1.kernel.shape}')
print(f'h.shape: {h.shape}')
layer.kernel.shape: (5, 8)
h.shape: (5, 8, 5, 8)

একটি এই চট ব্যবহার করার জন্য নিউটনের পদ্ধতি ধাপ, আপনাকে প্রথমে জানতে তার অক্ষ একটি ম্যাট্রিক্স মধ্যে চেপ্টা, এবং একটি ভেক্টর মধ্যে গ্রেডিয়েন্ট আউট চেপ্টা হবে:

n_params = tf.reduce_prod(layer1.kernel.shape)

g_vec = tf.reshape(g, [n_params, 1])
h_mat = tf.reshape(h, [n_params, n_params])

হেসিয়ান ম্যাট্রিক্স প্রতিসম হওয়া উচিত:

def imshow_zero_center(image, **kwargs):
  lim = tf.reduce_max(abs(image))
  plt.imshow(image, vmin=-lim, vmax=lim, cmap='seismic', **kwargs)
  plt.colorbar()
imshow_zero_center(h_mat)

png

নিউটনের পদ্ধতি আপডেট ধাপ নিচে দেখানো হয়েছে:

eps = 1e-3
eye_eps = tf.eye(h_mat.shape[0])*eps
# X(k+1) = X(k) - (∇²f(X(k)))^-1 @ ∇f(X(k))
# h_mat = ∇²f(X(k))
# g_vec = ∇f(X(k))
update = tf.linalg.solve(h_mat + eye_eps, g_vec)

# Reshape the update and apply it to the variable.
_ = layer1.kernel.assign_sub(tf.reshape(update, layer1.kernel.shape))

এই একটি একক জন্য অপেক্ষাকৃত সহজ tf.Variable , একটি অ-তুচ্ছ মডেল এই আবেদন সাবধান সংযুক্তকরণের এবং slicing প্রয়োজন একাধিক ভেরিয়েবল জুড়ে একটি পূর্ণ চট উত্পাদন করতে হবে।

ব্যাচ জ্যাকোবিয়ান

কিছু ক্ষেত্রে, আপনি উৎসের স্ট্যাকের সাপেক্ষে প্রতিটি লক্ষ্যের স্ট্যাকের জ্যাকোবিয়ান নিতে চান, যেখানে প্রতিটি লক্ষ্য-উৎস জোড়ার জন্য জ্যাকোবিয়ানরা স্বাধীন।

উদাহরণস্বরূপ, এখানে ইনপুট x চূড়ান্ত করা হয় (batch, ins) এবং আউটপুট y চূড়ান্ত করা হয় (batch, outs) :

x = tf.random.normal([7, 5])

layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)

with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
  tape.watch(x)
  y = layer1(x)
  y = layer2(y)

y.shape
TensorShape([7, 6])

পূর্ণ Jacobian y থেকে সম্মান সঙ্গে x একটি আকৃতির আছে (batch, ins, batch, outs) , এমনকি আপনি যদি শুধুমাত্র চান (batch, ins, outs) :

j = tape.jacobian(y, x)
j.shape
TensorShape([7, 6, 7, 5])

স্ট্যাকের মধ্যে প্রতিটি আইটেমের গ্রেডিয়েন্ট স্বাধীন, তাহলে প্রত্যেক (batch, batch) এই টেন্সর ফালি একটি তির্যক ম্যাট্রিক্স হল:

imshow_zero_center(j[:, 0, :, 0])
_ = plt.title('A (batch, batch) slice')

png

def plot_as_patches(j):
  # Reorder axes so the diagonals will each form a contiguous patch.
  j = tf.transpose(j, [1, 0, 3, 2])
  # Pad in between each patch.
  lim = tf.reduce_max(abs(j))
  j = tf.pad(j, [[0, 0], [1, 1], [0, 0], [1, 1]],
             constant_values=-lim)
  # Reshape to form a single image.
  s = j.shape
  j = tf.reshape(j, [s[0]*s[1], s[2]*s[3]])
  imshow_zero_center(j, extent=[-0.5, s[2]-0.5, s[0]-0.5, -0.5])

plot_as_patches(j)
_ = plt.title('All (batch, batch) slices are diagonal')

png

কাঙ্ক্ষিত ফলাফল পেতে, আপনি ডুপ্লিকেট উপর যোগফল করতে batch মাত্রা, বা অন্য ব্যবহার কর্ণ নির্বাচন tf.einsum :

j_sum = tf.reduce_sum(j, axis=2)
print(j_sum.shape)
j_select = tf.einsum('bxby->bxy', j)
print(j_select.shape)
(7, 6, 5)
(7, 6, 5)

প্রথম স্থানে অতিরিক্ত মাত্রা ছাড়াই গণনা করা অনেক বেশি কার্যকর হবে। tf.GradientTape.batch_jacobian পদ্ধতি ঠিক তাই করছে:

jb = tape.batch_jacobian(y, x)
jb.shape
WARNING:tensorflow:5 out of the last 5 calls to <function pfor.<locals>.f at 0x7f7d601250e0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
TensorShape([7, 6, 5])
error = tf.reduce_max(abs(jb - j_sum))
assert error < 1e-3
print(error.numpy())
0.0
x = tf.random.normal([7, 5])

layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)
bn = tf.keras.layers.BatchNormalization()
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)

with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
  tape.watch(x)
  y = layer1(x)
  y = bn(y, training=True)
  y = layer2(y)

j = tape.jacobian(y, x)
print(f'j.shape: {j.shape}')
WARNING:tensorflow:6 out of the last 6 calls to <function pfor.<locals>.f at 0x7f7cf062fa70> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
j.shape: (7, 6, 7, 5)
plot_as_patches(j)

_ = plt.title('These slices are not diagonal')
_ = plt.xlabel("Don't use `batch_jacobian`")

png

এই ক্ষেত্রে, batch_jacobian এখনো সঞ্চালিত হয় ও প্রত্যাশিত আকৃতি সঙ্গে আয় কিছু, কিন্তু এটা এর বিষয়বস্তু একটি অস্পষ্ট অর্থ আছে:

jb = tape.batch_jacobian(y, x)
print(f'jb.shape: {jb.shape}')
jb.shape: (7, 6, 5)