def train(self,
data_rdd,
model_rdd,
batch_size,
epochs,
model_dir,
go_on=False):
n_samples = data_rdd.count()
# steps_per_epoch = n_samples // batch_size // self.num_workers
steps_per_epoch = math.ceil(n_samples / batch_size / self.num_workers)
assert steps_per_epoch > 0
md = ModelDir(model_dir, 'train*')
if go_on:
md.create_model_dir()
else:
md = md.rebuild_model_dir()
worker = TFTrainWorker(model_rdd,
go_on=go_on,
batch_size=batch_size,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
**md.to_dict())
cluster = TFCluster.run(self.sc,
worker,
self.tf_args,
self.cluster_size,
self.num_ps,
input_mode=self.input_mode)
cluster.train(data_rdd.rdd, num_epochs=epochs, feed_timeout=60000)
cluster.shutdown()
results = md.read_result()
return self.sqlc.createDataFrame(results)
def main(self):
md = ModelDir(self.model_dir, 'train*')
if self.go_on:
md.create_model_dir()
else:
md = md.rebuild_model_dir()
self.build_model()
self.train(**md.to_dict())
def yolov3_tiny_train(self,
model_rdd,
batch_size,
epochs,
classes_path,
anchors_path,
train_path,
val_path,
image_size,
model_dir,
weights_path=None,
freeze_body=2,
go_on=False):
columns = model_rdd.columns
assert "model_config" in columns, "not exists model layer config!"
assert tf.io.gfile.exists(train_path), "train dataset path not exists!"
data_rdd = self.sc.textFile(train_path)
n_samples = data_rdd.count()
steps_per_epoch = math.ceil(n_samples / batch_size / self.num_workers)
md = ModelDir(model_dir, 'train*')
if go_on:
md.create_model_dir()
else:
md = md.rebuild_model_dir()
worker = YOLOV3TinyModelTrainWorker(model_rdd,
go_on=go_on,
batch_size=batch_size,
epochs=epochs,
classes_path=classes_path,
anchors_path=anchors_path,
weights_path=weights_path,
val_path=val_path,
image_size=image_size,
steps_per_epoch=steps_per_epoch,
freeze_body=freeze_body,
**md.to_dict())
cluster = TFCluster.run(self.sc,
worker,
self.tf_args,
self.cluster_size,
self.num_ps,
input_mode=self.input_mode)
cluster.train(data_rdd, num_epochs=epochs, feed_timeout=60000)
cluster.shutdown()
results = md.read_result()
return self.sqlc.createDataFrame(results)
def yolov3_train(self,
model_rdd,
data_dir,
batch_size,
epochs,
image_size,
model_dir,
weights_path=None,
freeze_body=2,
go_on=False):
train_path = os.path.join(data_dir, 'train.txt')
assert tf.io.gfile.exists(train_path), "train dataset path not exists!"
data_rdd = self.sc.textFile(train_path)
n_samples = data_rdd.count()
steps_per_epoch = math.ceil(n_samples / batch_size / self.num_workers)
md = ModelDir(model_dir, 'train*')
if go_on:
md.create_model_dir()
else:
md = md.rebuild_model_dir()
worker = YOLOV3ModelTrainWorker(model_rdd,
data_dir,
go_on=go_on,
batch_size=batch_size,
epochs=epochs,
image_size=image_size,
steps_per_epoch=steps_per_epoch,
freeze_body=freeze_body,
**md.to_dict())
cluster = TFCluster.run(self.sc,
worker,
self.tf_args,
self.cluster_size,
self.num_ps,
input_mode=self.input_mode)
cluster.train(data_rdd, num_epochs=epochs, feed_timeout=60000)
cluster.shutdown()
results = md.read_result()
if results:
return self.sqlc.createDataFrame(results)