| 301 | } |
| 302 | |
| 303 | func (m *TfkgModel) GetLayerWeights(layerName string) ([]*tf.Tensor, error) { |
| 304 | var variableOutputs []tf.Output |
| 305 | |
| 306 | found := false |
| 307 | for _, l := range m.layers { |
| 308 | if l.GetName() != layerName { |
| 309 | continue |
| 310 | } |
| 311 | found = true |
| 312 | for _, operation := range m.model.Graph.Operations() { |
| 313 | if strings.HasPrefix(operation.Name(), l.GetName()) && operation.Type() == "ReadVariableOp" { |
| 314 | variableOutputs = append(variableOutputs, m.model.Graph.Operation(operation.Name()).Output(0)) |
| 315 | } |
| 316 | } |
| 317 | } |
| 318 | if !found { |
| 319 | e := fmt.Errorf("layer %s not found in the model", layerName) |
| 320 | return nil, e |
| 321 | } |
| 322 | |
| 323 | results, e := m.model.Session.Run( |
| 324 | map[tf.Output]*tf.Tensor{}, |
| 325 | variableOutputs, |
| 326 | nil, |
| 327 | ) |
| 328 | if e != nil { |
| 329 | return nil, e |
| 330 | } |
| 331 | |
| 332 | return results, nil |
| 333 | } |
| 334 | |
| 335 | func (m *TfkgModel) SetModelWeights(weights []*tf.Tensor) error { |
| 336 | var variableOutputs []tf.Output |