-
Notifications
You must be signed in to change notification settings - Fork 29
/
custom_network_stub.py
60 lines (50 loc) · 1.72 KB
/
custom_network_stub.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
"""
Use a custom network.
"""
from adept.networks import NetworkModule, NetworkRegistry
from adept.scripts.local import parse_args, main
class MyCustomNetwork(NetworkModule):
# You will be prompted for these when training script starts
args = {"example_arg1": True, "example_arg2": 5}
def __init__(self):
super(MyCustomNetwork, self).__init__()
# Set properties and whatnot here
@classmethod
def from_args(cls, args, observation_space, output_space, net_reg):
"""
Construct a MyCustomNetwork from arguments.
ArgName = str
ObsKey = str
OutputKey = str
Shape = Tuple[*int]
:param args: Dict[ArgName, Any]
:param observation_space: Dict[ObsKey, Shape]
:param output_space: Dict[OutputKey, Shape]
:param net_reg: NetworkRegistry
:return: MyCustomNetwork
"""
pass
def new_internals(self, device):
"""
Define any initial hidden states here, move them to device if necessary.
InternalKey=str
:return: Dict[InternalKey, torch.Tensor (ND)]
"""
pass
def forward(self, observation, internals):
"""
Compute forward pass.
ObsKey = str
InternalKey = str
:param observation: Dict[ObsKey, torch.Tensor (1D | 2D | 3D | 4D)]
:param internals: Dict[InternalKey, torch.Tensor (ND)]
:return: torch.Tensor
"""
pass
if __name__ == "__main__":
args = parse_args()
network_reg = NetworkRegistry()
network_reg.register_network(MyCustomNetwork)
main(args, net_registry=network_reg)
# Call script like this to train agent:
# python -m custom_network_stub.py --custom-network MyCustomNetwork