Method Optimize

optimizers/optimizer.go:37–78

Optimize performs the optimization of the parameters.

()

// Optimize performs the optimization of the parameters.
func (o *Optimizer) Optimize() error {
	var wg sync.WaitGroup
	guard := make(chan struct{}, runtime.NumCPU()*2)
	errCh := make(chan error, 1)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	for param := range o.parameters(ctx) {
		select {
		case err := <-errCh:
			cancel()  // As soon as an error occurs, stop the iteration over parameters
			wg.Wait() // Wait for running goroutines to finish
			return err
		default:
			param := param
			wg.Add(1)
			guard <- struct{}{}
			go func() {
				defer wg.Done()
				defer func() { <-guard }()
				if !param.HasGrad() {
					return
				}
				if err := o.strategy.OptimizeParams(param); err != nil {
					select {
					case errCh <- err:
					default:
					}
				}
			}()
		}
	}

	close(errCh)

	if err, ok := <-errCh; ok {
		return err
	}

	return nil
}