50b78f02cfe965eceff936f739825f8eb981b8fb,CapsNet.py,,,#,66

Before Change

                  metrics={"out_caps": "accuracy"})

    // begin training
    model.fit([x_train, y_train], [y_train, x_train], batch_size=args.batch_size, epochs=args.epochs,
              validation_data=[[x_test, y_test], [y_test, x_test]])
    model.save("trained_model.h5")
    print("Trained model saved to \"trained_model.h5\"")

After Change

    // callbacks
    log = callbacks.CSVLogger(args.save_dir + "/log.csv")
    tb = callbacks.TensorBoard(log_dir=args.save_dir+"/tensorboard-logs", batch_size=args.batch_size)
    checkpoint = callbacks.ModelCheckpoint(args.save_dir + "/weights-{epoch:02d}.h5", save_best_only=True, verbose=1)

    // define model
    model = CapsNet(input_shape=[28, 28, 1],
                    n_class=len(np.unique(np.argmax(y_train, 1))),
                    batch_size=args.batch_size)
    model.summary()
    model.compile(optimizer="adam",
                  loss=[margin_loss, "mse"],
                  loss_weights=[1., args.lam_recon],
                  metrics={"out_caps": "accuracy"})

    // begin training
    // Training without data augmentation:
    // model.fit([x_train, y_train], [y_train, x_train], batch_size=args.batch_size, epochs=args.epochs,
    //           validation_data=[[x_test, y_test], [y_test, x_test]], callbacks=[log, tb])

    // Training with data augmentation. If shift_fraction=0., also no augmentation.
    model.fit_generator(generator=train_generator(x_train, y_train, args.batch_size, args.shift_fraction),
                        steps_per_epoch=int(y_train.shape[0] / args.batch_size),
                        epochs=args.epochs,
                        validation_data=[[x_test, y_test], [y_test, x_test]],
                        callbacks=[log, tb, checkpoint])
    model.save(args.save_dir + "/trained_model.h5")
    print("Trained model saved to \"trained_model.h5\"")

Instances