| 333 | } |
| 334 | |
| 335 | func (m *TfkgModel) SetModelWeights(weights []*tf.Tensor) error { |
| 336 | var variableOutputs []tf.Output |
| 337 | output := 0 |
| 338 | for _, info := range m.model.Signatures["set_weights"].Outputs { |
| 339 | parts := strings.Split(info.Name, ":") |
| 340 | if len(parts) != 2 { |
| 341 | e := fmt.Errorf("error getting output for set_weights signature in SetModelWeights") |
| 342 | m.errorHandler.Error(e) |
| 343 | return e |
| 344 | } |
| 345 | name := parts[0] |
| 346 | variableOutputs = append(variableOutputs, m.model.Graph.Operation(name).Output(output)) |
| 347 | output++ |
| 348 | } |
| 349 | var inputOps []tf.Output |
| 350 | for offset := range weights { |
| 351 | inputOps = append(inputOps, m.model.Graph.Operation(fmt.Sprintf("set_weights_weights_%d", offset)).Output(0)) |
| 352 | } |
| 353 | |
| 354 | inputs := map[tf.Output]*tf.Tensor{} |
| 355 | |
| 356 | for i, inputOp := range inputOps { |
| 357 | inputs[inputOp] = weights[i] |
| 358 | } |
| 359 | _, e := m.model.Session.Run( |
| 360 | inputs, |
| 361 | variableOutputs, |
| 362 | nil, |
| 363 | ) |
| 364 | if e != nil { |
| 365 | m.errorHandler.Error(e) |
| 366 | return e |
| 367 | } |
| 368 | |
| 369 | return nil |
| 370 | } |
| 371 | |
| 372 | func (m *TfkgModel) Predict(inputs ...*tf.Tensor) (*tf.Tensor, error) { |
| 373 | if len(inputs) < 1 { |