Image (NHWC) to patches (N(H' W')(patch_height patch_width c)).
(images, patch_height, patch_width)
| 22 | |
| 23 | |
| 24 | def to_patch(images, patch_height, patch_width): |
| 25 | """Image (NHWC) to patches (N(H' W')(patch_height patch_width c)).""" |
| 26 | batch_size, h, w, c = tf_utils.get_shape_list(images) |
| 27 | num_h = h // patch_height |
| 28 | num_w = w // patch_width |
| 29 | x = tf.reshape(images, |
| 30 | (batch_size, num_h, patch_height, num_w, patch_width, c)) |
| 31 | x = tf.einsum('nhpwqc->nhwpqc', x) |
| 32 | x = tf.reshape(x, (batch_size, num_h, num_w, patch_height * patch_width * c)) |
| 33 | return x |
| 34 | |
| 35 | |
| 36 | class ViTClassifier(tf_keras.Model): |