model.ChessNetFactorizedMoE Class Reference

Inheritance diagram for model.ChessNetFactorizedMoE:

Collaboration diagram for model.ChessNetFactorizedMoE:

Public Member Functions
	__init__ (self, n_globals=21, branch_dim=16, mixer_out=64, n_bypass=12, n_experts=4, expert_bottleneck=32, expert_hidden=128, expert_pool="flat")
	forward (self, *inputs, weights=None)

Public Attributes
	expert_pool = expert_pool
	n_experts = n_experts
	branches
	pointwise_mixer = nn.Conv2d(mixer_in_channels, mixer_out, kernel_size=1)
	stem_global = nn.Linear(n_globals, mixer_out)
	mixer_act = nn.ReLU(inplace=True)
	experts

Static Public Attributes
list	PLANES_PER_TYPE = [4, 4, 5, 5, 5, 4, 4, 4, 5, 5, 5, 4]

Constructor & Destructor Documentation

__init__()

model.ChessNetFactorizedMoE.__init__	(		self,
			n_globals = 21,
			branch_dim = 16,
			mixer_out = 64,
			n_bypass = 12,
			n_experts = 4,
			expert_bottleneck = 32,
			expert_hidden = 128,
			expert_pool = "flat" )

Here is the call graph for this function:

Here is the caller graph for this function:

Member Function Documentation

forward()

model.ChessNetFactorizedMoE.forward	(		self,
		*	inputs,
			weights = None )

Member Data Documentation

branches

model.ChessNetFactorizedMoE.branches

Initial value:

=  nn.ModuleList([
            PieceBranch(in_channels=in_ch, mid_channels=branch_dim) 
            for in_ch in self.PLANES_PER_TYPE
        ])

expert_pool

model.ChessNetFactorizedMoE.expert_pool = expert_pool

experts

model.ChessNetFactorizedMoE.experts

Initial value:

=  nn.ModuleList([
            LightExpert(
                mixer_channels=mixer_out, 
                bottleneck_channels=expert_bottleneck, 
                hidden_dim=expert_hidden,
                expert_pool=expert_pool,
            ) for _ in range(n_experts)
        ])

mixer_act

model.ChessNetFactorizedMoE.mixer_act = nn.ReLU(inplace=True)

n_experts

model.ChessNetFactorizedMoE.n_experts = n_experts

PLANES_PER_TYPE

list model.ChessNetFactorizedMoE.PLANES_PER_TYPE = [4, 4, 5, 5, 5, 4, 4, 4, 5, 5, 5, 4]

static

pointwise_mixer

model.ChessNetFactorizedMoE.pointwise_mixer = nn.Conv2d(mixer_in_channels, mixer_out, kernel_size=1)

stem_global

model.ChessNetFactorizedMoE.stem_global = nn.Linear(n_globals, mixer_out)

---

The documentation for this class was generated from the following file:

• training/model.py