Skip to content

Commit

Permalink
Orca pytorch estimator wrapper (intel-analytics#2624)
Browse files Browse the repository at this point in the history
* add estimator wrapper

* fix testcase

* fix style

* delete testcase & modify example
  • Loading branch information
cyita committed Jul 28, 2020
1 parent a8e5509 commit 5825485
Showing 1 changed file with 6 additions and 7 deletions.
13 changes: 6 additions & 7 deletions python/orca/example/learn/horovod/pytorch_estimator.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@

from zoo import init_spark_on_yarn, init_spark_on_local
from zoo.ray import RayContext
from zoo.orca.learn.pytorch.pytorch_horovod_estimator import PyTorchHorovodEstimator
from zoo.orca.learn.pytorch.estimator import Estimator


class LinearDataset(torch.utils.data.Dataset):
Expand Down Expand Up @@ -78,7 +78,7 @@ def validation_data_creator(config):


def train_example():
trainer1 = PyTorchHorovodEstimator(
estimator = Estimator.from_model_creator(
model_creator=model_creator,
optimizer_creator=optimizer_creator,
loss_creator=nn.MSELoss,
Expand All @@ -90,11 +90,10 @@ def train_example():
})

# train 5 epochs
for i in range(5):
stats = trainer1.train(train_data_creator)
print("train stats: {}".format(stats))
val_stats = trainer1.validate(validation_data_creator)
print("validation stats: {}".format(val_stats))
stats = estimator.fit(train_data_creator, epochs=5)
print("train stats: {}".format(stats))
val_stats = estimator.evaluate(validation_data_creator)
print("validation stats: {}".format(val_stats))


parser = argparse.ArgumentParser()
Expand Down

0 comments on commit 5825485

Please sign in to comment.