| 370 | } |
| 371 | |
| 372 | func (m *TfkgModel) Predict(inputs ...*tf.Tensor) (*tf.Tensor, error) { |
| 373 | if len(inputs) < 1 { |
| 374 | e := fmt.Errorf("no inputs provided") |
| 375 | m.errorHandler.Error(e) |
| 376 | return nil, e |
| 377 | } |
| 378 | var predictOutputs []tf.Output |
| 379 | output := 0 |
| 380 | for _, info := range m.model.Signatures["predict"].Outputs { |
| 381 | parts := strings.Split(info.Name, ":") |
| 382 | if len(parts) != 2 { |
| 383 | e := fmt.Errorf("error getting output for predict signature in fit") |
| 384 | m.errorHandler.Error(e) |
| 385 | return nil, e |
| 386 | } |
| 387 | name := parts[0] |
| 388 | predictOutputs = append(predictOutputs, m.model.Graph.Operation(name).Output(output)) |
| 389 | output++ |
| 390 | } |
| 391 | predictInputs := map[tf.Output]*tf.Tensor{} |
| 392 | for i, inputTensor := range inputs { |
| 393 | predictInputs[m.model.Graph.Operation(fmt.Sprintf("predict_inputs_%d", i)).Output(0)] = inputTensor |
| 394 | } |
| 395 | |
| 396 | results, e := m.model.Session.Run( |
| 397 | predictInputs, |
| 398 | predictOutputs, |
| 399 | nil, |
| 400 | ) |
| 401 | if e != nil { |
| 402 | m.errorHandler.Error(e) |
| 403 | return nil, e |
| 404 | } |
| 405 | |
| 406 | return results[0], nil |
| 407 | } |
| 408 | |
| 409 | type FitConfig struct { |
| 410 | Epochs int |