def train_mnist(args): track.init(trial_name="track-example", trial_config=vars(args)) batch_size = 128 num_classes = 10 epochs = 1 if args.smoke_test else 12 mnist.load() x_train, y_train, x_test, y_test, input_shape = get_mnist_data() model = Sequential() model.add( Conv2D(32, kernel_size=(3, 3), activation="relu", input_shape=input_shape)) model.add(Conv2D(64, (3, 3), activation="relu")) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Dropout(0.5)) model.add(Flatten()) model.add(Dense(args.hidden, activation="relu")) model.add(Dropout(0.5)) model.add(Dense(num_classes, activation="softmax")) model.compile(loss="categorical_crossentropy", optimizer=keras.optimizers.SGD(lr=args.lr, momentum=args.momentum), metrics=["accuracy"]) model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_data=(x_test, y_test), callbacks=[TuneKerasCallback(track.metric)]) track.shutdown()
def train_mnist(args): track.init(trial_name="track-example", trial_config=vars(args)) batch_size = 128 num_classes = 10 epochs = 1 if args.smoke_test else 12 (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train, x_test = x_train / 255.0, x_test / 255.0 model = tf.keras.models.Sequential([ tf.keras.layers.Flatten(input_shape=(28, 28)), tf.keras.layers.Dense(args.hidden, activation="relu"), tf.keras.layers.Dropout(0.2), tf.keras.layers.Dense(num_classes, activation="softmax") ]) model.compile(loss="sparse_categorical_crossentropy", optimizer=keras.optimizers.SGD(lr=args.lr, momentum=args.momentum), metrics=["accuracy"]) model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_data=(x_test, y_test), callbacks=[TuneReporterCallback()]) track.shutdown()
def testSessionInitShutdown(self): self.assertTrue(track._session is None) # Checks that the singleton _session is created/destroyed # by track.init() and track.shutdown() for _ in range(2): # do it twice to see that we can reopen the session track.init(trial_name="test_init") self.assertTrue(track._session is not None) track.shutdown() self.assertTrue(track._session is None)
def _trainable_func(self, config, reporter): track.init(_tune_reporter=reporter) output = train_func(config) reporter(**{RESULT_DUPLICATE: True}) track.shutdown() return output
def tearDown(self): track.shutdown() ray.shutdown()