ऑटो-बैच संयुक्त वितरण: एक कोमल ट्यूटोरियल

TensorFlow.org पर देखें Google Colab में चलाएं GitHub पर स्रोत देखें नोटबुक डाउनलोड करें

परिचय

TensorFlow संभावना (टीएफपी) के एक नंबर प्रदान करता है JointDistribution आसानी से करने के लिए एक उपयोगकर्ता एक लगभग गणितीय रूप में एक संभाव्य चित्रमय मॉडल व्यक्त की अनुमति देकर आसान कपोल-कल्पना उस संभावित निष्कर्ष बनाते हैं; अमूर्त मॉडल से नमूने लेने और मॉडल से नमूनों की लॉग संभावना का मूल्यांकन करने के लिए तरीके उत्पन्न करता है। इस ट्यूटोरियल में, हम की समीक्षा "autobatched" वेरिएंट, जो मूल के बाद विकसित किए गए JointDistribution कपोल-कल्पना। मूल, गैर-ऑटोबैच्ड एब्स्ट्रैक्शन के सापेक्ष, ऑटोबैच्ड संस्करण उपयोग में आसान और अधिक एर्गोनोमिक हैं, जिससे कई मॉडल कम बॉयलरप्लेट के साथ व्यक्त किए जा सकते हैं। इस कोलाब में, हम (शायद थकाऊ) विस्तार में एक सरल मॉडल का पता लगाते हैं, जिससे ऑटोबैचिंग हल करने वाली समस्याओं को स्पष्ट करते हैं, और (उम्मीद है) पाठक को रास्ते में टीएफपी आकार अवधारणाओं के बारे में अधिक सिखाते हैं।

Autobatching की शुरूआत करने से पहले, वहाँ के कुछ अलग-अलग रूपों थे JointDistribution , संभाव्य मॉडल व्यक्त करने के लिए अलग वाक्यात्मक शैलियों के लिए इसी: JointDistributionSequential , JointDistributionNamed , और JointDistributionCoroutine । Auobatching एक mixin के रूप में मौजूद है, इसलिए अब हमारे पास AutoBatched इन सब के वेरिएंट। इस ट्यूटोरियल में, हम के बीच मतभेद का पता लगाने JointDistributionSequential और JointDistributionSequentialAutoBatched ; हालांकि, हम यहां जो कुछ भी करते हैं वह अन्य वेरिएंट पर लागू होता है जिसमें अनिवार्य रूप से कोई बदलाव नहीं होता है।

निर्भरता और पूर्वापेक्षाएँ

आयात और सेट अप

पूर्वापेक्षा: एक बायेसियन रिग्रेशन समस्या

हम एक बहुत ही सरल बायेसियन प्रतिगमन परिदृश्य पर विचार करेंगे:

\[ \begin{align*} m & \sim \text{Normal}(0, 1) \\ b & \sim \text{Normal}(0, 1) \\ Y & \sim \text{Normal}(mX + b, 1) \end{align*} \]

इस मॉडल में, m और b मानक normals से लिए गए हैं, और टिप्पणियों Y एक सामान्य वितरण जिसका मतलब यादृच्छिक परिवर्तनीय पर निर्भर करता है से तैयार कर रहे हैं m और b , और कुछ (nonrandom, जाना जाता है) covariates X । (सादगी के लिए, इस उदाहरण में, हम मानते हैं कि सभी यादृच्छिक चरों का पैमाना ज्ञात है।)

इस मॉडल में निष्कर्ष करने के लिए, हम दोनों covariates पता करने की जरूरत होगी X और टिप्पणियों Y , लेकिन इस ट्यूटोरियल के प्रयोजनों के लिए, हम केवल आवश्यकता होगी X , तो हम एक साधारण डमी परिभाषित X :

X = np.arange(7)
X
array([0, 1, 2, 3, 4, 5, 6])

देसीराता

संभाव्य अनुमान में, हम अक्सर दो बुनियादी संचालन करना चाहते हैं:

  • sample : मॉडल से ड्राइंग नमूने हैं।
  • log_prob : मॉडल से एक नमूना की लॉग संभावना गणना की जा रही।

TFP के मुख्य योगदान JointDistribution कपोल-कल्पना (और साथ ही की संभाव्य प्रोग्रामिंग करने के लिए कई अन्य तरीकों) उन एक बार एक मॉडल लिख सकते हैं और दोनों का उपयोग करने की अनुमति देने के लिए है sample और log_prob संगणना।

ध्यान देने योग्य बात है कि हम अपने डेटा सेट (7 अंक है X.shape = (7,) ), अब हम एक उत्कृष्ट के लिए Desiderata कह सकते हैं JointDistribution :

  • sample() की एक सूची तैयार करना चाहिए Tensors आकार वाले [(), (), (7,) ] अदिश ढलान, अदिश पूर्वाग्रह, और वेक्टर टिप्पणियों के लिए इसी क्रमश।
  • log_prob(sample()) एक विशेष ढाल, पूर्वाग्रह, और टिप्पणियों की लॉग संभावना: एक अदिश प्रस्तुत करना चाहिए।
  • sample([5, 3]) की एक सूची तैयार करना चाहिए Tensors आकार वाले [(5, 3), (5, 3), (5, 3, 7)] , एक का प्रतिनिधित्व (5, 3) - से नमूने के बैच आदर्श।
  • log_prob(sample([5, 3])) एक प्रस्तुत करना चाहिए Tensor आकार (5, 3) के साथ।

अब हम का एक उत्तराधिकार को देखेंगे JointDistribution , मॉडल कैसे ऊपर Desiderata प्राप्त करने के लिए देखते हैं, और उम्मीद है कि एक छोटे से अधिक जानने के बारे में TFP रास्ते आकार देती है।

स्पॉइलर चेतावनी: दृष्टिकोण है कि संतुष्ट जोड़ा बॉयलरप्लेट बिना ऊपर Desiderata है autobatching

पहला प्रयास; JointDistributionSequential

jds = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Normal(loc=m*X + b, scale=1.) # Y
])

यह कमोबेश मॉडल का कोड में सीधा अनुवाद है। ढलान m और पूर्वाग्रह b सीधा कर रहे हैं। Y एक का उपयोग कर परिभाषित किया गया है lambda समारोह: सामान्य पैटर्न है कि एक है lambda के समारोह \(k\) एक में तर्क JointDistributionSequential (जेडीएस) पिछले उपयोग करता \(k\) मॉडल में वितरण। "रिवर्स" ऑर्डर पर ध्यान दें।

हम फोन करता हूँ sample_distributions , जो रिटर्न दोनों एक नमूना और अंतर्निहित "उप वितरण" है कि नमूना उत्पन्न करने के लिए इस्तेमाल किया गया। (हम फोन करके सिर्फ नमूना उत्पादन कर सकते थे sample : नमूना हम उत्पादन ठीक है, साथ ही बाद में ट्यूटोरियल में यह वितरण के लिए सुविधाजनक हो जाएगा।)

dists, sample = jds.sample_distributions()
sample
[<tf.Tensor: shape=(), dtype=float32, numpy=-1.668757>,
 <tf.Tensor: shape=(), dtype=float32, numpy=0.6585061>,
 <tf.Tensor: shape=(7,), dtype=float32, numpy=
 array([ 0.18573815, -1.79962   , -1.8106272 , -3.5971394 , -6.6625295 ,
        -7.308844  , -9.832693  ], dtype=float32)>]

लेकिन log_prob किसी अवांछित आकार के साथ एक परिणाम पैदा करता है:

jds.log_prob(sample)
<tf.Tensor: shape=(7,), dtype=float32, numpy=
array([-4.4777603, -4.6775575, -4.7430477, -4.647725 , -4.5746684,
       -4.4368567, -4.480562 ], dtype=float32)>

और एकाधिक नमूनाकरण काम नहीं करता है:

try:
  jds.sample([5, 3])
except tf.errors.InvalidArgumentError as e:
  print(e)
Incompatible shapes: [5,3] vs. [7] [Op:Mul]

आइए समझने की कोशिश करें कि क्या गलत हो रहा है।

एक संक्षिप्त समीक्षा: बैच और घटना आकार

TFP में, एक साधारण (एक नहीं JointDistribution ) प्रायिकता वितरण एक घटना आकार और एक बैच आकार, और अंतर को समझने TFP के प्रभावी उपयोग के लिए महत्वपूर्ण है है:

  • घटना आकार वितरण से एकल ड्रा के आकार का वर्णन करता है; ड्रा सभी आयामों पर निर्भर हो सकता है। अदिश वितरण के लिए, घटना का आकार [] है। 5-आयामी बहुभिन्नरूपी सामान्य के लिए, घटना का आकार [5] है।
  • बैच आकार स्वतंत्र का वर्णन करता है, समान रूप से वितरित ड्रॉ नहीं, उर्फ ​​​​वितरण का "बैच"। एकल पायथन ऑब्जेक्ट में वितरण के एक बैच का प्रतिनिधित्व करना उन प्रमुख तरीकों में से एक है जिनसे TFP बड़े पैमाने पर दक्षता प्राप्त करता है।

हमारे प्रयोजनों के लिए, मन में रखने के लिए एक महत्वपूर्ण तथ्य यह है कि अगर हम कहते हैं log_prob एक वितरण से एक भी नमूना पर, परिणाम हमेशा एक आकार है कि मैचों (यानी, सबसे दायीं ओर आयाम के रूप में है) बैच आकार होगा।

आकार का एक और अधिक में गहराई से चर्चा के लिए, देखें "समझौता TensorFlow वितरण आकृतियाँ" ट्यूटोरियल

क्यों नहीं log_prob(sample()) एक अदिश निर्माण?

चलो बैच और घटना आकार के हमारे ज्ञान का उपयोग क्या साथ क्या हो रहा है पता लगाने के लिए log_prob(sample()) यहाँ हमारा नमूना फिर से है:

sample
[<tf.Tensor: shape=(), dtype=float32, numpy=-1.668757>,
 <tf.Tensor: shape=(), dtype=float32, numpy=0.6585061>,
 <tf.Tensor: shape=(7,), dtype=float32, numpy=
 array([ 0.18573815, -1.79962   , -1.8106272 , -3.5971394 , -6.6625295 ,
        -7.308844  , -9.832693  ], dtype=float32)>]

और यहाँ हमारे वितरण हैं:

dists
[<tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>,
 <tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>,
 <tfp.distributions.Normal 'JointDistributionSequential_sample_distributions_Normal' batch_shape=[7] event_shape=[] dtype=float32>]

लॉग संभाव्यता की गणना भागों के (मिलान) तत्वों पर उप-वितरण की लॉग संभावनाओं को जोड़कर की जाती है:

log_prob_parts = [dist.log_prob(s) for (dist, s) in zip(dists, sample)]
log_prob_parts
[<tf.Tensor: shape=(), dtype=float32, numpy=-2.3113134>,
 <tf.Tensor: shape=(), dtype=float32, numpy=-1.1357536>,
 <tf.Tensor: shape=(7,), dtype=float32, numpy=
 array([-1.0306933, -1.2304904, -1.2959809, -1.200658 , -1.1276014,
        -0.9897899, -1.0334952], dtype=float32)>]
np.sum(log_prob_parts) - jds.log_prob(sample)
<tf.Tensor: shape=(7,), dtype=float32, numpy=array([0., 0., 0., 0., 0., 0., 0.], dtype=float32)>

तो, विवरण में से एक स्तर है कि लॉग संभावना गणना एक 7 टेन्सर लौटा रहा है क्योंकि के तीसरे उप-घटक है log_prob_parts एक 7 टेन्सर है। लेकिन क्यों?

ठीक है, हम देखते हैं कि के अंतिम तत्व dists , जिस पर अपने वितरण से मेल खाती है Y mathematial निर्माण में, एक है batch_shape की [7] । दूसरे शब्दों में, हमारे वितरण Y (इस मामले में, एक ही पैमाने में विभिन्न साधनों के साथ और,) 7 स्वतंत्र normals का एक बैच है।

अब हम समझते हैं कि क्या गलत है: जेडीएस में, से अधिक वितरण Y है batch_shape=[7] , जेडीएस से एक नमूना के लिए scalars का प्रतिनिधित्व करता है m और b और एक 7 स्वतंत्र normals की "बैच"। और log_prob 7 अलग लॉग-संभावनाओं, जिनमें से प्रत्येक ड्राइंग का लॉग संभावना दर्शाते गणना करता m और b और एक भी अवलोकन Y[i] कुछ पर X[i]

फिक्सिंग log_prob(sample()) के साथ Independent

याद रखें कि dists[2] है event_shape=[] और batch_shape=[7] :

dists[2]
<tfp.distributions.Normal 'JointDistributionSequential_sample_distributions_Normal' batch_shape=[7] event_shape=[] dtype=float32>

TFP के का उपयोग करके Independent metadistribution है, जो घटना आयामों को बैच आयाम बदल देता है, हम के साथ एक वितरण में इस परिवर्तित कर सकते हैं event_shape=[7] और batch_shape=[] (हम इसका नाम बदलना करेंगे y_dist_i क्योंकि उस पर एक समान वितरण है Y , साथ _i खड़े हमारे लिए में Independent रैपिंग):

y_dist_i = tfd.Independent(dists[2], reinterpreted_batch_ndims=1)
y_dist_i
<tfp.distributions.Independent 'IndependentJointDistributionSequential_sample_distributions_Normal' batch_shape=[] event_shape=[7] dtype=float32>

अब, log_prob एक 7 वेक्टर के एक अदिश है:

y_dist_i.log_prob(sample[2])
<tf.Tensor: shape=(), dtype=float32, numpy=-7.9087086>

के तहत कवर, Independent बैच से अधिक रकम:

y_dist_i.log_prob(sample[2]) - tf.reduce_sum(dists[2].log_prob(sample[2]))
<tf.Tensor: shape=(), dtype=float32, numpy=0.0>

और वास्तव में, हम एक नया निर्माण करने के लिए इसका उपयोग कर सकते jds_i ( i फिर से के लिए खड़ा है Independent ) जहां log_prob एक अदिश रिटर्न:

jds_i = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Independent(   # Y
        tfd.Normal(loc=m*X + b, scale=1.),
        reinterpreted_batch_ndims=1)
])

jds_i.log_prob(sample)
<tf.Tensor: shape=(), dtype=float32, numpy=-11.355776>

एक युगल नोट:

  • jds_i.log_prob(s) के समान नहीं है tf.reduce_sum(jds.log_prob(s)) पूर्व संयुक्त वितरण की "सही" लॉग संभावना पैदा करता है। एक 7 टेन्सर से अधिक बाद रकम, प्रत्येक तत्व जिनमें से लॉग संभावना का योग है m , b , और लॉग संभावना की एक एकल तत्व Y , तो यह overcounts m और b । ( log_prob(m) + log_prob(b) + log_prob(Y) एक अपवाद फेंक क्योंकि TFP TF और NumPy के प्रसारण नियमों का पालन करती बजाय एक परिणाम देता है;। एक वेक्टर के लिए एक अदिश जोड़ने एक वेक्टर आकार परिणाम पैदा करता है)
  • इस विशेष मामले में, हम समस्या को हल किया जा सकता था और का उपयोग कर एक ही परिणाम हासिल MultivariateNormalDiag के बजाय Independent(Normal(...)) MultivariateNormalDiag एक वेक्टर-मान वितरण है (यानी, यह पहले से ही वेक्टर घटना आकार है)। Indeeed MultivariateNormalDiag हो सकता है (लेकिन नहीं है) की एक रचना के रूप में लागू Independent और Normal । ऐसा नहीं है कि एक सदिश दी याद करने के लिए फ़ायदेमंद हो सकता है V , से नमूने n1 = Normal(loc=V) , और n2 = MultivariateNormalDiag(loc=V) अप्रभेद्य हैं; ये वितरण beween अंतर यह है कि है n1.log_prob(n1.sample()) एक वेक्टर और है n2.log_prob(n2.sample()) एक अदिश है।

एकाधिक नमूने?

कई नमूने खींचना अभी भी काम नहीं करता है:

try:
  jds_i.sample([5, 3])
except tf.errors.InvalidArgumentError as e:
  print(e)
Incompatible shapes: [5,3] vs. [7] [Op:Mul]

आइए विचार करें कि क्यों। जब हम कहते हैं jds_i.sample([5, 3]) , हम पहले नमूने के लिए आकर्षित करेंगे m और b , आकार के साथ प्रत्येक (5, 3) । इसके बाद, हम एक निर्माण करने के लिए प्रयास करने के लिए जा रहे हैं Normal के माध्यम से वितरण:

tfd.Normal(loc=m*X + b, scale=1.)

लेकिन अगर m आकार (5, 3) और X आकार 7 , हम गुणा उन्हें एक साथ नहीं कर सकते हैं, और वास्तव में इस त्रुटि हम हिटिंग है:

m = tfd.Normal(0., 1.).sample([5, 3])
try:
  m * X
except tf.errors.InvalidArgumentError as e:
  print(e)
Incompatible shapes: [5,3] vs. [7] [Op:Mul]

इस समस्या को हल करने के लिए, के गुणों से अधिक वितरण क्या बारे में सोचते हैं Y है है। हमें बुलाया गया है jds_i.sample([5, 3]) , तो हम जानते हैं कि m और b दोनों आकार होगा (5, 3) । क्या आकार एक कॉल करने के लिए करना चाहिए sample पर Y वितरण उपज? स्पष्ट जवाब है (5, 3, 7) : प्रत्येक बैच बिंदु के लिए, हम के रूप में एक ही आकार के साथ एक नमूना चाहते X । हम अतिरिक्त आयाम जोड़कर, TensorFlow की प्रसारण क्षमताओं का उपयोग करके इसे प्राप्त कर सकते हैं:

m[..., tf.newaxis].shape
TensorShape([5, 3, 1])
(m[..., tf.newaxis] * X).shape
TensorShape([5, 3, 7])

दोनों के लिए एक धुरी जोड़ना m और b , हम एक नए जेडीएस परिभाषित कर सकते हैं कि समर्थन करता है कई नमूने:

jds_ia = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Independent(   # Y
        tfd.Normal(loc=m[..., tf.newaxis]*X + b[..., tf.newaxis], scale=1.),
        reinterpreted_batch_ndims=1)
])

shaped_sample = jds_ia.sample([5, 3])
shaped_sample
[<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
 array([[-1.1133379 ,  0.16390413, -0.24177533],
        [-1.1312429 , -0.6224666 , -1.8182136 ],
        [-0.31343174, -0.32932565,  0.5164407 ],
        [-0.0119963 , -0.9079621 ,  2.3655841 ],
        [-0.26293617,  0.8229698 ,  0.31098196]], dtype=float32)>,
 <tf.Tensor: shape=(5, 3), dtype=float32, numpy=
 array([[-0.02876974,  1.0872147 ,  1.0138507 ],
        [ 0.27367726, -1.331534  , -0.09084719],
        [ 1.3349475 , -0.68765205,  1.680652  ],
        [ 0.75436825,  1.3050154 , -0.9415123 ],
        [-1.2502679 , -0.25730947,  0.74611956]], dtype=float32)>,
 <tf.Tensor: shape=(5, 3, 7), dtype=float32, numpy=
 array([[[-1.8258233e+00, -3.0641669e-01, -2.7595463e+00, -1.6952467e+00,
          -4.8197951e+00, -5.2986512e+00, -6.6931367e+00],
         [ 3.6438566e-01,  1.0067395e+00,  1.4542470e+00,  8.1155670e-01,
           1.8868095e+00,  2.3877139e+00,  1.0195159e+00],
         [-8.3624744e-01,  1.2518480e+00,  1.0943471e+00,  1.3052304e+00,
          -4.5756745e-01, -1.0668410e-01, -7.0669651e-02]],
 
        [[-3.1788960e-01,  9.2615485e-03, -3.0963073e+00, -2.2846246e+00,
          -3.2269263e+00, -6.0213070e+00, -7.4806519e+00],
         [-3.9149747e+00, -3.5155020e+00, -1.5669601e+00, -5.0759468e+00,
          -4.5065498e+00, -5.6719379e+00, -4.8012795e+00],
         [ 1.3053948e-01, -8.0493152e-01, -4.7845001e+00, -4.9721808e+00,
          -7.1365709e+00, -9.6198196e+00, -9.7951422e+00]],
 
        [[ 2.0621397e+00,  3.4639853e-01,  7.0252883e-01, -1.4311566e+00,
           3.3790007e+00,  1.1619035e+00, -8.9105040e-01],
         [-7.8956139e-01, -8.5023916e-01, -9.7148323e-01, -2.6229355e+00,
          -2.7150445e+00, -2.4633870e+00, -2.1841538e+00],
         [ 7.7627432e-01,  2.2401071e+00,  3.7601702e+00,  2.4245868e+00,
           4.0690269e+00,  4.0605016e+00,  5.1753912e+00]],
 
        [[ 1.4275590e+00,  3.3346462e+00,  1.5374103e+00, -2.2849756e-01,
           9.1219616e-01, -3.1220305e-01, -3.2643962e-01],
         [-3.1910419e-02, -3.8848895e-01,  9.9946201e-02, -2.3619974e+00,
          -1.8507402e+00, -3.6830821e+00, -5.4907336e+00],
         [-7.1941972e-02,  2.1602919e+00,  4.9575748e+00,  4.2317696e+00,
           9.3528280e+00,  1.0526063e+01,  1.5262107e+01]],
 
        [[-2.3257759e+00, -2.5343289e+00, -3.5342445e+00, -4.0423255e+00,
          -3.2361765e+00, -3.3434000e+00, -2.6849220e+00],
         [ 1.5006512e-02, -1.9866472e-01,  7.6781356e-01,  1.6228745e+00,
           1.4191239e+00,  2.6655579e+00,  4.4663467e+00],
         [ 2.6599693e+00,  1.2663836e+00,  1.7162113e+00,  1.4839669e+00,
           2.0559487e+00,  2.5976877e+00,  2.5977583e+00]]], dtype=float32)>]
jds_ia.log_prob(shaped_sample)
<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[-12.483114 , -10.139662 , -11.514159 ],
       [-11.656767 , -17.201958 , -12.132455 ],
       [-17.838818 ,  -9.474525 , -11.24898  ],
       [-13.95219  , -12.490049 , -17.123957 ],
       [-14.487818 , -11.3755455, -10.576363 ]], dtype=float32)>

एक अतिरिक्त जांच के रूप में, हम सत्यापित करेंगे कि एकल बैच बिंदु के लिए लॉग संभावना हमारे पास पहले की तुलना में मेल खाती है:

(jds_ia.log_prob(shaped_sample)[3, 1] -
 jds_i.log_prob([shaped_sample[0][3, 1],
                 shaped_sample[1][3, 1],
                 shaped_sample[2][3, 1, :]]))
<tf.Tensor: shape=(), dtype=float32, numpy=0.0>

जीत के लिए ऑटोबैचिंग

उत्कृष्ट! अब हम JointDistribution का एक संस्करण है कि हैंडल सब हमारे Desiderata: log_prob रिटर्न का उपयोग करने के एक अदिश धन्यवाद tfd.Independent , और कई नमूने अब काम है कि हम अतिरिक्त कुल्हाड़ियों जोड़कर प्रसारण तय की।

क्या होगा अगर मैंने आपको बताया कि एक आसान, बेहतर तरीका था? वहाँ है, और यह कहा जाता है JointDistributionSequentialAutoBatched (JDSAB):

jds_ab = tfd.JointDistributionSequentialAutoBatched([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Normal(loc=m*X + b, scale=1.) # Y
])
jds_ab.log_prob(jds.sample())
<tf.Tensor: shape=(), dtype=float32, numpy=-12.954952>
shaped_sample = jds_ab.sample([5, 3])
jds_ab.log_prob(shaped_sample)
<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[-12.191533 , -10.43885  , -16.371655 ],
       [-13.292994 , -11.97949  , -16.788685 ],
       [-15.987699 , -13.435732 , -10.6029   ],
       [-10.184758 , -11.969714 , -14.275676 ],
       [-12.740775 , -11.5654125, -12.990162 ]], dtype=float32)>
jds_ab.log_prob(shaped_sample) - jds_ia.log_prob(shaped_sample)
<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float32)>

यह कैसे काम करता है? आप करने का प्रयास कर सकता है जबकि कोड को पढ़ने के लिए एक गहरी समझ के लिए, हम एक संक्षिप्त सिंहावलोकन जो सबसे उपयोग के मामलों के लिए पर्याप्त है प्रदान करेंगे:

  • याद रखें कि हमारी पहली समस्या यह है कि के लिए अपने वितरण था Y था batch_shape=[7] और event_shape=[] , और हम इस्तेमाल किया Independent एक घटना आयाम करने के लिए बैच आयाम कन्वर्ट करने के लिए। JDSAB घटक वितरण के बैच आकार की उपेक्षा करता है; बजाय यह मॉडल है, जो माना जाता है की एक समग्र संपत्ति के रूप में बैच आकार व्यवहार करता है [] (जब तक की स्थापना द्वारा अन्यथा निर्दिष्ट batch_ndims > 0 )। प्रभाव tfd.Independent का उपयोग कर घटना आयामों में घटक वितरण के सभी बैच आयाम परिवर्तित करने के लिए, के रूप में हम स्वयं ऊपर किया के बराबर है।
  • हमारी दूसरी समस्या के आकार की मालिश की आवश्यकता हुई m और b इतना है कि वे साथ उचित रूप से प्रसारित कर सकता है X जब कई नमूने का निर्माण। JDSAB के साथ, आप एक मॉडल एक भी नमूना उत्पन्न करने के लिए लिखते हैं, और हम "उठाना" पूरे मॉडल TensorFlow के का उपयोग कर कई नमूने उत्पन्न करने के लिए vectorized_map । (यह सुविधा JAX के लिए analagous है vmap ।)

और अधिक विस्तार में बैच आकार मुद्दा तलाश, हम अपने मूल "बुरा" संयुक्त वितरण के बैच आकार की तुलना कर सकते jds , हमारे बैच निर्धारित वितरण jds_i और jds_ia , और हमारे autobatched jds_ab :

jds.batch_shape
[TensorShape([]), TensorShape([]), TensorShape([7])]
jds_i.batch_shape
[TensorShape([]), TensorShape([]), TensorShape([])]
jds_ia.batch_shape
[TensorShape([]), TensorShape([]), TensorShape([])]
jds_ab.batch_shape
TensorShape([])

हम देखते हैं कि मूल jds अलग बैच आकार के साथ subdistributions है। jds_i और jds_ia ही (खाली) बैच आकार के साथ subdistributions बनाने के द्वारा इसे ठीक। jds_ab केवल एक ही (खाली) बैच आकार है।

यह लायक यह देखते हुए कि JointDistributionSequentialAutoBatched मुक्त करने के लिए कुछ अतिरिक्त व्यापकता प्रदान करता है। हम covariates बनाने मान लीजिए X (और, परोक्ष, टिप्पणियों Y ) दो आयामी:

X = np.arange(14).reshape((2, 7))
X
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13]])

हमारे JointDistributionSequentialAutoBatched कोई बदलाव नहीं (हम मॉडल को फिर से परिभाषित करने के लिए, क्योंकि के आकार की जरूरत के साथ काम करता X द्वारा कैश किया गया है jds_ab.log_prob ):

jds_ab = tfd.JointDistributionSequentialAutoBatched([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Normal(loc=m*X + b, scale=1.) # Y
])

shaped_sample = jds_ab.sample([5, 3])
shaped_sample
[<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
 array([[ 0.1813647 , -0.85994506,  0.27593774],
        [-0.73323774,  1.1153806 ,  0.8841938 ],
        [ 0.5127983 , -0.29271227,  0.63733214],
        [ 0.2362284 , -0.919168  ,  1.6648189 ],
        [ 0.26317367,  0.73077047,  2.5395133 ]], dtype=float32)>,
 <tf.Tensor: shape=(5, 3), dtype=float32, numpy=
 array([[ 0.09636458,  2.0138032 , -0.5054413 ],
        [ 0.63941646, -1.0785882 , -0.6442188 ],
        [ 1.2310615 , -0.3293852 ,  0.77637213],
        [ 1.2115169 , -0.98906034, -0.07816773],
        [-1.1318136 ,  0.510014  ,  1.036522  ]], dtype=float32)>,
 <tf.Tensor: shape=(5, 3, 2, 7), dtype=float32, numpy=
 array([[[[-1.9685398e+00, -1.6832136e+00, -6.9127172e-01,
            8.5992378e-01, -5.3123581e-01,  3.1584005e+00,
            2.9044402e+00],
          [-2.5645006e-01,  3.1554163e-01,  3.1186538e+00,
            1.4272424e+00,  1.2843871e+00,  1.2266440e+00,
            1.2798605e+00]],
 
         [[ 1.5973477e+00, -5.3631151e-01,  6.8143606e-03,
           -1.4910895e+00, -2.1568544e+00, -2.0513713e+00,
           -3.1663666e+00],
          [-4.9448099e+00, -2.8385928e+00, -6.9027486e+00,
           -5.6543546e+00, -7.2378774e+00, -8.1577444e+00,
           -9.3582869e+00]],
 
         [[-2.1233239e+00,  5.8853775e-02,  1.2024102e+00,
            1.6622503e+00, -1.9197327e-01,  1.8647723e+00,
            6.4322817e-01],
          [ 3.7549341e-01,  1.5853541e+00,  2.4594500e+00,
            2.1952972e+00,  1.7517658e+00,  2.9666045e+00,
            2.5468128e+00]]],
 
 
        [[[ 8.9906776e-01,  6.7375046e-01,  7.3354661e-01,
           -9.9894643e-01, -3.4606690e+00, -3.4810467e+00,
           -4.4315586e+00],
          [-3.0670738e+00, -6.3628020e+00, -6.2538433e+00,
           -6.8091092e+00, -7.7134805e+00, -8.6319380e+00,
           -8.6904278e+00]],
 
         [[-2.2462025e+00, -3.3060855e-01,  1.8974400e-01,
            3.1422038e+00,  4.1483402e+00,  3.5642972e+00,
            4.8709240e+00],
          [ 4.7880130e+00,  5.8790064e+00,  9.6695948e+00,
            7.8112822e+00,  1.2022618e+01,  1.2411858e+01,
            1.4323385e+01]],
 
         [[-1.0189297e+00, -7.8115642e-01,  1.6466728e+00,
            8.2378983e-01,  3.0765080e+00,  3.0170646e+00,
            5.1899948e+00],
          [ 6.5285158e+00,  7.8038850e+00,  6.4155884e+00,
            9.0899811e+00,  1.0040427e+01,  9.1404457e+00,
            1.0411951e+01]]],
 
 
        [[[ 4.5557004e-01,  1.4905317e+00,  1.4904103e+00,
            2.9777462e+00,  2.8620450e+00,  3.4745665e+00,
            3.8295493e+00],
          [ 3.9977460e+00,  5.7173767e+00,  7.8421035e+00,
            6.3180594e+00,  6.0838981e+00,  8.2257290e+00,
            9.6548376e+00]],
 
         [[-7.0750320e-01, -3.5972297e-01,  4.3136525e-01,
           -2.3301599e+00, -5.0374687e-01, -2.8338656e+00,
           -3.4453444e+00],
          [-3.1258626e+00, -3.4687450e+00, -1.2045374e+00,
           -4.0196013e+00, -5.8831010e+00, -4.2965469e+00,
           -4.1388311e+00]],
 
         [[ 2.1969774e+00,  2.4614549e+00,  2.2314475e+00,
            1.8392437e+00,  2.8367062e+00,  4.8600502e+00,
            4.2273531e+00],
          [ 6.1879644e+00,  5.1792760e+00,  6.1141996e+00,
            5.6517797e+00,  8.9979610e+00,  7.5938139e+00,
            9.7918644e+00]]],
 
 
        [[[ 1.5249090e+00,  1.1388919e+00,  8.6903995e-01,
            3.0762129e+00,  1.5128503e+00,  3.5204377e+00,
            2.4760864e+00],
          [ 3.4166217e+00,  3.5930209e+00,  3.1694956e+00,
            4.5797420e+00,  4.5271711e+00,  2.8774328e+00,
            4.7288942e+00]],
 
         [[-2.3095846e+00, -2.0595703e+00, -3.0093951e+00,
           -3.8594103e+00, -4.9681158e+00, -6.4256043e+00,
           -5.5345035e+00],
          [-6.4306297e+00, -7.0924540e+00, -8.4075985e+00,
           -1.0417805e+01, -1.1727266e+01, -1.1196255e+01,
           -1.1333830e+01]],
 
         [[-7.0419472e-01,  1.4568675e+00,  3.7946482e+00,
            4.8489718e+00,  6.6498446e+00,  9.0224218e+00,
            1.1153137e+01],
          [ 1.0060651e+01,  1.1998097e+01,  1.5326431e+01,
            1.7957514e+01,  1.8323889e+01,  2.0160881e+01,
            2.1269085e+01]]],
 
 
        [[[-2.2360647e-01, -1.3632748e+00, -7.2704530e-01,
            2.3558271e-01, -1.0381399e+00,  1.9387857e+00,
           -3.3694571e-01],
          [ 1.6015106e-01,  1.5284677e+00, -4.8567140e-01,
           -1.7770648e-01,  2.1919653e+00,  1.3015286e+00,
            1.3877077e+00]],
 
         [[ 1.3688663e+00,  2.6602898e+00,  6.6657305e-01,
            4.6554832e+00,  5.7781887e+00,  4.9115267e+00,
            4.8446012e+00],
          [ 5.1983776e+00,  6.2297459e+00,  6.3848300e+00,
            8.4291229e+00,  7.1309576e+00,  1.0395646e+01,
            8.5736713e+00]],
 
         [[ 1.2675294e+00,  5.2844582e+00,  5.1331611e+00,
            8.9993315e+00,  1.0794343e+01,  1.4039831e+01,
            1.5731170e+01],
          [ 1.9084715e+01,  2.2191265e+01,  2.3481146e+01,
            2.5803375e+01,  2.8632090e+01,  3.0234968e+01,
            3.1886738e+01]]]], dtype=float32)>]
jds_ab.log_prob(shaped_sample)
<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[-28.90071 , -23.052422, -19.851362],
       [-19.775568, -25.894997, -20.302256],
       [-21.10754 , -23.667885, -20.973007],
       [-19.249458, -20.87892 , -20.573763],
       [-22.351208, -25.457762, -24.648403]], dtype=float32)>

दूसरी ओर, हमारे सावधानी से गढ़ी गई JointDistributionSequential नहीं रह गया काम करता है:

jds_ia = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Independent(   # Y
        tfd.Normal(loc=m[..., tf.newaxis]*X + b[..., tf.newaxis], scale=1.),
        reinterpreted_batch_ndims=1)
])

try:
  jds_ia.sample([5, 3])
except tf.errors.InvalidArgumentError as e:
  print(e)
Incompatible shapes: [5,3,1] vs. [2,7] [Op:Mul]

इसे ठीक करने के लिए, हम एक दूसरे को जोड़ने के लिए होगा tf.newaxis दोनों को m और b आकार, और वृद्धि से मेल reinterpreted_batch_ndims करने के लिए कॉल में से 2 Independent । इस मामले में, ऑटो-बैचिंग मशीनरी को आकार के मुद्दों को संभालने देना छोटा, आसान और अधिक एर्गोनोमिक है।

एक बार फिर, हम ध्यान दें कि जब इस नोटबुक पता लगाया JointDistributionSequentialAutoBatched , की अन्य विविधताओं JointDistribution बराबर AutoBatched । (के उपयोगकर्ताओं के लिए JointDistributionCoroutine , JointDistributionCoroutineAutoBatched अतिरिक्त लाभ है कि आप निर्दिष्ट करने के लिए आवश्यकता नहीं Root नोड्स, यदि आप उपयोग नहीं किया है JointDistributionCoroutine । आप सुरक्षित रूप से इस बयान को नजरअंदाज कर सकते हैं)

समापन विचार

इस नोटबुक में, हम शुरू की JointDistributionSequentialAutoBatched और विस्तार से एक सरल उदाहरण के माध्यम से काम किया। उम्मीद है कि आपने TFP आकृतियों और ऑटोबैचिंग के बारे में कुछ सीखा होगा!