Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
6 contributors

Users who have contributed to this file

@asimshankar @jhseu @frreiss @jtkeeling @tensorflower-gardener @anight
409 lines (365 sloc) 11.8 KB
/*
Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package tensorflow
// #include <stdlib.h>
// #include "tensorflow/c/c_api.h"
import "C"
import (
"errors"
"fmt"
"runtime"
"sync"
"unsafe"
)
// Session drives a TensorFlow graph computation.
//
// When a Session is created with a given target, a new Session object is bound
// to the universe of resources specified by that target. Those resources are
// available to this session to perform computation described in the GraphDef.
// After creating the session with a graph, the caller uses the Run() API to
// perform the computation and potentially fetch outputs as Tensors.
// A Session allows concurrent calls to Run().
type Session struct {
c *C.TF_Session
// For ensuring that:
// - Close() blocks on all Run() calls to complete.
// - Close() can be called multiple times.
wg sync.WaitGroup
mu sync.Mutex
}
// NewSession creates a new execution session with the associated graph.
// options may be nil to use the default options.
func NewSession(graph *Graph, options *SessionOptions) (*Session, error) {
status := newStatus()
cOpt, doneOpt, err := options.c()
defer doneOpt()
if err != nil {
return nil, err
}
cSess := C.TF_NewSession(graph.c, cOpt, status.c)
if err := status.Err(); err != nil {
return nil, err
}
s := &Session{c: cSess}
runtime.SetFinalizer(s, func(s *Session) { s.Close() })
return s, nil
}
// Device structure contains information about a device associated with a session, as returned by ListDevices()
type Device struct {
Name, Type string
MemoryLimitBytes int64
}
// String describes d and implements fmt.Stringer.
func (d Device) String() string {
memStr := "no memory limit"
if d.MemoryLimitBytes >= 0 {
memStr = fmt.Sprintf("memory limit %d bytes", d.MemoryLimitBytes)
}
return fmt.Sprintf("(Device: name \"%s\", type %s, %s)", d.Name, d.Type, memStr)
}
func deviceSliceFromDeviceList(list *C.TF_DeviceList) ([]Device, error) {
var devices []Device
status := newStatus()
for i := 0; i < int(C.TF_DeviceListCount(list)); i++ {
name := C.TF_DeviceListName(list, C.int(i), status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("DeviceListName(index=%d) failed: %v", i, err)
}
deviceType := C.TF_DeviceListType(list, C.int(i), status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("DeviceListType(index=%d) failed: %v", i, err)
}
memoryLimitBytes := C.TF_DeviceListMemoryBytes(list, C.int(i), status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("DeviceListMemoryBytes(index=%d) failed: %v", i, err)
}
device := Device{
Name: C.GoString(name),
Type: C.GoString(deviceType),
MemoryLimitBytes: int64(memoryLimitBytes),
}
devices = append(devices, device)
}
return devices, nil
}
// ListDevices returns the list of devices associated with a Session.
func (s *Session) ListDevices() ([]Device, error) {
status := newStatus()
devicesList := C.TF_SessionListDevices(s.c, status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("SessionListDevices() failed: %v", err)
}
defer C.TF_DeleteDeviceList(devicesList)
return deviceSliceFromDeviceList(devicesList)
}
// Run the graph with the associated session starting with the supplied feeds
// to compute the value of the requested fetches. Runs, but does not return
// Tensors for operations specified in targets.
//
// On success, returns the fetched Tensors in the same order as supplied in
// the fetches argument. If fetches is set to nil, the returned Tensor fetches
// is empty.
func (s *Session) Run(feeds map[Output]*Tensor, fetches []Output, targets []*Operation) ([]*Tensor, error) {
s.mu.Lock()
if s.c == nil {
s.mu.Unlock()
return nil, errors.New("session is closed")
}
s.wg.Add(1)
s.mu.Unlock()
defer s.wg.Done()
c := newCRunArgs(feeds, fetches, targets)
status := newStatus()
C.TF_SessionRun(s.c, nil,
ptrOutput(c.feeds), ptrTensor(c.feedTensors), C.int(len(feeds)),
ptrOutput(c.fetches), ptrTensor(c.fetchTensors), C.int(len(fetches)),
ptrOperation(c.targets), C.int(len(targets)),
nil, status.c)
// Make sure GC won't harvest input tensors until SessionRun() is finished
runtime.KeepAlive(feeds)
if err := status.Err(); err != nil {
return nil, err
}
return c.toGo(), nil
}
// PartialRun enables incremental evaluation of graphs.
//
// PartialRun allows the caller to pause the evaluation of a graph, run
// arbitrary code that depends on the intermediate computation of the graph,
// and then resume graph execution. The results of the arbitrary code can be
// fed into the graph when resuming execution. In contrast, Session.Run
// executes the graph to compute the requested fetches using the provided feeds
// and discards all intermediate state (e.g., value of intermediate tensors)
// when it returns.
//
// For example, consider a graph for unsupervised training of a neural network
// model. PartialRun can be used to pause execution after the forward pass of
// the network, let the caller actuate the output (e.g., play a game, actuate a
// robot etc.), determine the error/loss and then feed this calculated loss
// when resuming the backward pass of the graph.
type PartialRun struct {
session *Session
handle *C.char
}
// Run resumes execution of the graph to compute the requested fetches and
// targets with the provided feeds.
func (pr *PartialRun) Run(feeds map[Output]*Tensor, fetches []Output, targets []*Operation) ([]*Tensor, error) {
var (
c = newCRunArgs(feeds, fetches, targets)
status = newStatus()
s = pr.session
)
s.mu.Lock()
if s.c == nil {
s.mu.Unlock()
return nil, errors.New("session is closed")
}
s.wg.Add(1)
s.mu.Unlock()
defer s.wg.Done()
C.TF_SessionPRun(s.c, pr.handle,
ptrOutput(c.feeds), ptrTensor(c.feedTensors), C.int(len(feeds)),
ptrOutput(c.fetches), ptrTensor(c.fetchTensors), C.int(len(fetches)),
ptrOperation(c.targets), C.int(len(targets)),
status.c)
if err := status.Err(); err != nil {
return nil, err
}
return c.toGo(), nil
}
// NewPartialRun sets up the graph for incremental evaluation.
//
// All values of feeds, fetches and targets that may be provided to Run calls
// on the returned PartialRun need to be provided to NewPartialRun.
//
// See documentation for the PartialRun type.
func (s *Session) NewPartialRun(feeds, fetches []Output, targets []*Operation) (*PartialRun, error) {
var (
cfeeds = make([]C.TF_Output, len(feeds))
cfetches = make([]C.TF_Output, len(fetches))
ctargets = make([]*C.TF_Operation, len(targets))
pcfeeds *C.TF_Output
pcfetches *C.TF_Output
pctargets **C.TF_Operation
status = newStatus()
)
if len(feeds) > 0 {
pcfeeds = &cfeeds[0]
for i, o := range feeds {
cfeeds[i] = o.c()
}
}
if len(fetches) > 0 {
pcfetches = &cfetches[0]
for i, o := range fetches {
cfetches[i] = o.c()
}
}
if len(targets) > 0 {
pctargets = &ctargets[0]
for i, o := range targets {
ctargets[i] = o.c
}
}
s.mu.Lock()
if s.c == nil {
s.mu.Unlock()
return nil, errors.New("session is closed")
}
s.wg.Add(1)
s.mu.Unlock()
defer s.wg.Done()
pr := &PartialRun{session: s}
C.TF_SessionPRunSetup(s.c,
pcfeeds, C.int(len(feeds)),
pcfetches, C.int(len(fetches)),
pctargets, C.int(len(targets)),
&pr.handle, status.c)
if err := status.Err(); err != nil {
return nil, err
}
runtime.SetFinalizer(pr, func(pr *PartialRun) {
C.TF_DeletePRunHandle(pr.handle)
})
return pr, nil
}
// Close a session. This contacts any other processes associated with this
// session, if applicable. Blocks until all previous calls to Run have returned.
func (s *Session) Close() error {
s.mu.Lock()
defer s.mu.Unlock()
s.wg.Wait()
if s.c == nil {
return nil
}
status := newStatus()
C.TF_CloseSession(s.c, status.c)
if err := status.Err(); err != nil {
return err
}
C.TF_DeleteSession(s.c, status.c)
s.c = nil
return status.Err()
}
// SessionOptions contains configuration information for a session.
type SessionOptions struct {
// Target indicates the TensorFlow runtime to connect to.
//
// If 'target' is empty or unspecified, the local TensorFlow runtime
// implementation will be used. Otherwise, the TensorFlow engine
// defined by 'target' will be used to perform all computations.
//
// "target" can be either a single entry or a comma separated list
// of entries. Each entry is a resolvable address of one of the
// following formats:
// local
// ip:port
// host:port
// ... other system-specific formats to identify tasks and jobs ...
//
// NOTE: at the moment 'local' maps to an in-process service-based
// runtime.
//
// Upon creation, a single session affines itself to one of the
// remote processes, with possible load balancing choices when the
// "target" resolves to a list of possible processes.
//
// If the session disconnects from the remote process during its
// lifetime, session calls may fail immediately.
Target string
// Config is a binary-serialized representation of the
// tensorflow.ConfigProto protocol message
// (https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto).
Config []byte
}
// c converts the SessionOptions to the C API's TF_SessionOptions. Callers must
// deallocate by calling the returned done() closure.
func (o *SessionOptions) c() (ret *C.TF_SessionOptions, done func(), err error) {
opt := C.TF_NewSessionOptions()
if o == nil {
return opt, func() { C.TF_DeleteSessionOptions(opt) }, nil
}
t := C.CString(o.Target)
C.TF_SetTarget(opt, t)
C.free(unsafe.Pointer(t))
var cConfig unsafe.Pointer
if sz := len(o.Config); sz > 0 {
status := newStatus()
// Copying into C-memory is the simplest thing to do in terms
// of memory safety and cgo rules ("C code may not keep a copy
// of a Go pointer after the call returns" from
// https://golang.org/cmd/cgo/#hdr-Passing_pointers).
cConfig = C.CBytes(o.Config)
C.TF_SetConfig(opt, cConfig, C.size_t(sz), status.c)
if err := status.Err(); err != nil {
C.TF_DeleteSessionOptions(opt)
return nil, func() {}, fmt.Errorf("invalid SessionOptions.Config: %v", err)
}
}
return opt, func() {
C.TF_DeleteSessionOptions(opt)
C.free(cConfig)
}, nil
}
// cRunArgs translates the arguments to Session.Run and PartialRun.Run into
// values suitable for C library calls.
type cRunArgs struct {
feeds []C.TF_Output
feedTensors []*C.TF_Tensor
fetches []C.TF_Output
fetchTensors []*C.TF_Tensor
targets []*C.TF_Operation
}
func newCRunArgs(feeds map[Output]*Tensor, fetches []Output, targets []*Operation) *cRunArgs {
c := &cRunArgs{
fetches: make([]C.TF_Output, len(fetches)),
fetchTensors: make([]*C.TF_Tensor, len(fetches)),
targets: make([]*C.TF_Operation, len(targets)),
}
for o, t := range feeds {
c.feeds = append(c.feeds, o.c())
c.feedTensors = append(c.feedTensors, t.c)
}
for i, o := range fetches {
c.fetches[i] = o.c()
}
for i, t := range targets {
c.targets[i] = t.c
}
return c
}
func (c *cRunArgs) toGo() []*Tensor {
ret := make([]*Tensor, len(c.fetchTensors))
for i, ct := range c.fetchTensors {
ret[i] = newTensorFromC(ct)
}
return ret
}
func ptrOutput(l []C.TF_Output) *C.TF_Output {
if len(l) == 0 {
return nil
}
return &l[0]
}
func ptrTensor(l []*C.TF_Tensor) **C.TF_Tensor {
if len(l) == 0 {
return nil
}
return &l[0]
}
func ptrOperation(l []*C.TF_Operation) **C.TF_Operation {
if len(l) == 0 {
return nil
}
return &l[0]
}
You can’t perform that action at this time.