| 21 | self.with_routing = with_routing |
| 22 | |
| 23 | def __call__(self, input, num_outputs, kernel_size=None, stride=None): |
| 24 | self.num_outputs = num_outputs |
| 25 | self.kernel_size = kernel_size |
| 26 | self.stride = stride |
| 27 | |
| 28 | if not self.with_routing: |
| 29 | # the PrimaryCaps layer |
| 30 | # input: [batch_size, 20, 20, 256] |
| 31 | assert input.get_shape() == [cfg.batch_size, 20, 20, 256] |
| 32 | |
| 33 | capsules = [] |
| 34 | for i in range(self.num_units): |
| 35 | # each capsule i: [batch_size, 6, 6, 32] |
| 36 | with tf.variable_scope('ConvUnit_' + str(i)): |
| 37 | caps_i = tf.contrib.layers.conv2d(input, |
| 38 | self.num_outputs, |
| 39 | self.kernel_size, |
| 40 | self.stride, |
| 41 | padding="VALID") |
| 42 | caps_i = tf.reshape(caps_i, shape=(cfg.batch_size, -1, 1, 1)) |
| 43 | capsules.append(caps_i) |
| 44 | |
| 45 | assert capsules[0].get_shape() == [cfg.batch_size, 1152, 1, 1] |
| 46 | |
| 47 | # [batch_size, 1152, 8, 1] |
| 48 | capsules = tf.concat(capsules, axis=2) |
| 49 | capsules = squash(capsules) |
| 50 | assert capsules.get_shape() == [cfg.batch_size, 1152, 8, 1] |
| 51 | |
| 52 | else: |
| 53 | # the DigitCaps layer |
| 54 | # Reshape the input into shape [batch_size, 1152, 8, 1] |
| 55 | self.input = tf.reshape(input, shape=(cfg.batch_size, 1152, 8, 1)) |
| 56 | |
| 57 | # b_IJ: [1, num_caps_l, num_caps_l_plus_1, 1] |
| 58 | b_IJ = tf.zeros(shape=[1, 1152, 10, 1], dtype=np.float32) |
| 59 | capsules = [] |
| 60 | for j in range(self.num_outputs): |
| 61 | with tf.variable_scope('caps_' + str(j)): |
| 62 | caps_j, b_IJ = capsule(input, b_IJ, j) |
| 63 | capsules.append(caps_j) |
| 64 | |
| 65 | # Return a tensor with shape [batch_size, 10, 16, 1] |
| 66 | capsules = tf.concat(capsules, axis=1) |
| 67 | assert capsules.get_shape() == [cfg.batch_size, 10, 16, 1] |
| 68 | |
| 69 | return(capsules) |
| 70 | |
| 71 | |
| 72 | def capsule(input, b_IJ, idx_j): |