Function train

test/integration/data/elastic_tensorflow2_main.py:118–151

(state)

@hvd.elastic.run
def train(state):
    state.rendezvous += 1
    while state.epoch < args.epochs:
        print('epoch {} batch {}'.format(state.epoch, state.batch))

        while state.batch < args.batches_per_epoch:
            check_exit(state.epoch, state.batch)
            step()

            state.batch += 1
            if state.batch % args.batches_per_commit == 0:
                state.commits += 1
                state.commit()

        if hvd.rank() == 0:
            log_state(state)

            current_hosts = epoch_to_hosts.get(state.epoch, default_hosts)
            next_hosts = epoch_to_hosts.get(state.epoch + 1, default_hosts)
            if args.discovery_wait > 0 and current_hosts != next_hosts:
                print('host changes: {} -> {}'.format(current_hosts, next_hosts))
                start = int(time.time())
                while state._host_messages.empty():
                    if int(time.time()) - start > args.discovery_wait:
                        raise TimeoutError('Timed out waiting for notifications from driver.')
                    time.sleep(0.1)

        if args.epoch_wait > 0:
            time.sleep(args.epoch_wait)

        state.epoch += 1
        state.batch = 0
        state.commits += 1
        state.commit()


def on_state_reset():