trX = np.expand_dims(trX[:, :, :, 1], axis=3)
vaX = np.expand_dims(vaX[:, :, :, 1], axis=3)
teX = np.expand_dims(teX[:, :, :, 1], axis=3)
EndTime = time.time()
cprint('Takes ' + str(EndTime - StartTime) + 'time to load data',
'magenta')
print(str(trX.shape), str(trY.shape), str(vaX.shape), str(vaY.shape))
# Define Model
#print(trY)
#cprint(str(np.argmax(trY,1)),'red')
cprint(str(len(np.unique(np.argmax(trY, 1)))), 'red')
cprint(str(len(np.unique(np.argmax(vaY, 1)))), 'red')
with tf.device('/cpu:0'):
model, eval_model, manipulate_model = CapsNet(
input_shape=trX.shape[1:],
n_class=len(np.unique(np.argmax(trY, 1))),
kernel=args.kernel,
primary_channel=args.primary_channel,
routings=args.routings)
#model = CapsNet_NoDecoder(input_shape=trX.shape[1:],
# n_class=len(np.unique(np.argmax(trY, 1))),
# kernel=args.kernel,
# primary_channel=args.primary_channel,
# routings=args.routings)
model.summary()
multi_model = multi_gpu_model(model, gpus=args.gpus)
# Save path and load model
if not os.path.exists(save_path):
os.mkdir(save_path)
if args.keep: # init the model weights with provided one
cprint('load weight from:' + save_path, 'yellow')
multi_model.load_weights(save_path)
print(args)
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# load data
if args.dataset == 0:
(x_train, y_train), (x_test, y_test) = load_mnist()
elif args.dataset == 1:
(x_train, y_train), (x_test, y_test) = load_fashion_mnist()
elif args.dataset == 2:
(x_train, y_train), (x_test, y_test) = load_svhn()
# define model
model, eval_model, manipulate_model = CapsNet(
input_shape=x_train.shape[1:],
n_class=len(np.unique(np.argmax(y_train, 1))),
routings=args.routings,
l1=args.l1)
model.summary()
flags = [0] * 10
index = [0] * 10
digits = np.where(y_test == 1)[1]
for i, num in enumerate(digits):
num = int(num)
if flags[num]:
continue
else:
flags[num] = 1
index[num] = i
if np.all(flags):
parser.add_argument('--is_training', default=1, type=int)
parser.add_argument('--weights', default=None)
parser.add_argument('--lr', default=0.001, type=float)
parser.add_argument('--gpus', default=2, type=int)
args = parser.parse_args()
print(args)
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# load data
(x_train, y_train), (x_test, y_test) = load_mnist()
# define model
with tf.device('/cpu:0'):
model, eval_model = CapsNet(input_shape=x_train.shape[1:],
n_class=len(
np.unique(np.argmax(y_train, 1))),
num_routing=args.num_routing)
model.summary()
plot_model(model, to_file=args.save_dir + '/model.png', show_shapes=True)
# define muti-gpu model
multi_model = multi_gpu_model(model, gpus=args.gpus)
# train or test
if args.weights is not None: # init the model weights with provided one
model.load_weights(args.weights)
if args.is_training:
train(model=multi_model,
data=((x_train, y_train), (x_test, y_test)),
args=args)
model.save_weights(args.save_dir + '/trained_model.h5')
print('Trained model saved to \'%s/trained_model.h5\'' % args.save_dir)
def _build_model(self):
self.model, self.target_model = CapsNet(self.state_size, self.action_size, 3)
data = DATA(args.is_training,
args.train_with,
args.test_with,
args.data_path,
mode=args.mode) #[sample,99,40,3]
X = data[0]
Y = data[1]
# Define Model
cprint(str(len(np.unique(np.argmax(Y, 1)))), 'red')
with tf.device('/cpu:0'):
if args.decoder == 1:
model, eval_model, manipulate_model = CapsNet(
input_shape=X.shape[1:],
n_class=len(np.unique(np.argmax(Y, 1))),
kernel=args.kernel,
primary_channel=args.primary_channel,
primary_veclen=args.primary_veclen,
digit_veclen=args.digit_veclen,
dropout=args.dropout,
routings=args.routings)
else:
model = CapsNet_NoDecoder(input_shape=X.shape[1:],
n_class=len(np.unique(np.argmax(Y, 1))),
kernel=args.kernel,
primary_channel=args.primary_channel,
primary_veclen=args.primary_veclen,
digit_veclen=args.digit_veclen,
dropout=args.dropout,
routings=args.routings)
model.summary()
multi_model = multi_gpu_model(model, gpus=args.gpus)