Struct WithCpuCopy

pub struct WithCpuCopy<T: AsModule> { /* private fields */ }

Wraps a module to have a lazily-initialized CPU copy if not already in CPU memory.

This is useful for models used both in training and in simulation because large batch size training is most efficient on the GPU while batch-size-1 simulation is most efficient on the CPU.

Implementations

impl<T: AsModule> WithCpuCopy<T>

pub const fn new(inner: T, device: Device) -> Self

impl<T: AsModule> WithCpuCopy<T>

pub const fn as_inner(&self) -> &T

Get a reference to the inner struct.

pub fn as_inner_mut(&mut self) -> &mut T

Get a mutable reference to the inner struct. Invalidates the cached CPU copy if any.

pub fn shallow_clone_module_cpu(&self) -> T::Module

Create a shallow clone of the module on CPU memory.

If the module is not already on the CPU device then a cached deep copy is created on the CPU first. This cached copy is reused on future calls.

Trait Implementations

impl<T: AsModule> AsModule for WithCpuCopy<T>

fn as_module_mut(&mut self) -> &mut Self::Module

Get mutable reference to the module. Invalidates the cached CPU copy if any.

type Module = <T as AsModule>::Module

fn as_module(&self) -> &Self::Module

fn batch_map<F: Fn(Tensor) -> Tensor>(self, f: F) -> BatchMap<Self, F>

where Self: Sized,

Apply a batch-structure-preserving function to the module output.

impl<T: AsModule + Clone> Clone for WithCpuCopy<T>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<T: AsModule + Debug> Debug for WithCpuCopy<T>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<'de, T> Deserialize<'de> for WithCpuCopy<T>

where T: Deserialize<'de> + AsModule,

fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl<T: AsModule + PartialEq> PartialEq for WithCpuCopy<T>

fn eq(&self, other: &Self) -> bool

Tests for self and other values to be equal, and is used by ==.

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl<P: Policy> Policy for WithCpuCopy<P>

type PolicyModule = <P as Policy>::PolicyModule

fn update<AS: ParameterizedDistributionSpace<Tensor> + ?Sized>( &mut self, features: &dyn HistoryFeatures, advantages: PackedTensor, action_space: &AS, logger: &mut dyn StatsLogger, )

Update the policy module.

fn actor<OS, AS>( &self, observation_space: NonEmptyFeatures<OS>, action_space: AS, ) -> PolicyActor<OS, AS, Self::Module>

Create an actor for the policy module.

impl<T> Serialize for WithCpuCopy<T>

where T: Serialize + AsModule,

fn serialize<__S>(&self, __serializer: __S) -> Result<__S::Ok, __S::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.