diff --git a/src/capsules/nets.py b/src/capsules/nets.py index f70c783..98b98f9 100644 --- a/src/capsules/nets.py +++ b/src/capsules/nets.py @@ -3,7 +3,7 @@ import tensorflow as tf -from core import _conv2d_wrapper, capsules_init, capsules_conv, capsules_fc +from capsules.core import _conv2d_wrapper, capsules_init, capsules_conv, capsules_fc slim = tf.contrib.slim @@ -26,6 +26,10 @@ def capsules_v0(inputs, num_classes, iterations, name='CapsuleEM-V0'): nets, shape=[1, 1, 32, 32], strides=[1, 1, 1, 1], padding='VALID', pose_shape=[4, 4], name='capsule_init' ) # inputs: (poses, activations) -> capsule-conv 3x3x32x32x4x4, strides 2 -> (poses, activations) + nets = capsules_conv( + nets, shape=[3, 3, 32, 32], strides=[1, 3, 3, 1], iterations=iterations, name='capsule_conv0' + ) + nets = capsules_conv( nets, shape=[3, 3, 32, 32], strides=[1, 2, 2, 1], iterations=iterations, name='capsule_conv1' ) diff --git a/src/settings.py b/src/settings.py index b782c08..3749068 100644 --- a/src/settings.py +++ b/src/settings.py @@ -4,7 +4,7 @@ from __future__ import division from __future__ import print_function -import capsules +from capsules import __init__ import tensorflow as tf diff --git a/src/train.py b/src/train.py index b2b40cb..f5cc10b 100644 --- a/src/train.py +++ b/src/train.py @@ -1,6 +1,6 @@ """A train script for matrix capsule with EM routing.""" -from settings import * +from . import settings def main(_):