-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathpytorch_util.py
More file actions
executable file
·187 lines (132 loc) · 4.7 KB
/
Copy pathpytorch_util.py
File metadata and controls
executable file
·187 lines (132 loc) · 4.7 KB
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
import torch
import numpy as np
def soft_update_from_to(source, target, tau):
for target_param, param in zip(target.parameters(), source.parameters()):
target_param.data.copy_(
target_param.data * (1.0 - tau) + param.data * tau
)
def copy_model_params_from_to(source, target):
for target_param, param in zip(target.parameters(), source.parameters()):
target_param.data.copy_(param.data)
def fanin_init(tensor, scale=1):
size = tensor.size()
if len(size) == 2:
fan_in = size[0]
elif len(size) > 2:
fan_in = np.prod(size[1:])
else:
raise Exception("Shape must be have dimension at least 2.")
bound = scale / np.sqrt(fan_in)
return tensor.data.uniform_(-bound, bound)
def orthogonal_init(tensor, gain=0.01):
torch.nn.init.orthogonal_(tensor, gain=gain)
def fanin_init_weights_like(tensor):
size = tensor.size()
if len(size) == 2:
fan_in = size[0]
elif len(size) > 2:
fan_in = np.prod(size[1:])
else:
raise Exception("Shape must be have dimension at least 2.")
bound = 1. / np.sqrt(fan_in)
new_tensor = FloatTensor(tensor.size())
new_tensor.uniform_(-bound, bound)
return new_tensor
"""
GPU wrappers
"""
_use_gpu = False
device = None
_gpu_id = 0
def set_gpu_mode(mode, gpu_id=0):
global _use_gpu
global device
global _gpu_id
_gpu_id = gpu_id
_use_gpu = mode
device = torch.device("cuda:" + str(gpu_id) if _use_gpu else "cpu")
def gpu_enabled():
return _use_gpu
def set_device(gpu_id):
torch.cuda.set_device(gpu_id)
# noinspection PyPep8Naming
def FloatTensor(*args, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.FloatTensor(*args, **kwargs, device=torch_device)
def from_numpy(*args, **kwargs):
return torch.from_numpy(*args, **kwargs).float().to(device)
def get_numpy(tensor):
return tensor.to('cpu').detach().numpy()
def zeros(*sizes, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.zeros(*sizes, **kwargs, device=torch_device)
def ones(*sizes, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.ones(*sizes, **kwargs, device=torch_device)
def ones_like(*args, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.ones_like(*args, **kwargs, device=torch_device)
def rand(*args, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.rand(*args, **kwargs, device=torch_device)
def randn(*args, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.randn(*args, **kwargs, device=torch_device)
def zeros_like(*args, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.zeros_like(*args, **kwargs, device=torch_device)
def tensor(*args, torch_device=None, **kwargs):
if torch_device is None:
torch_device = device
return torch.tensor(*args, **kwargs, device=torch_device)
def normal(*args, **kwargs):
return torch.normal(*args, **kwargs).to(device)
def eval_np(module, *args, **kwargs):
"""
Eval this module with a numpy interface
Same as a call to __call__ except all Variable input/outputs are
replaced with numpy equivalents.
Assumes the output is either a single object or a tuple of objects.
"""
torch_args = tuple(torch_ify(x) for x in args)
torch_kwargs = {k: torch_ify(v) for k, v in kwargs.items()}
outputs = module(*torch_args, **torch_kwargs)
if isinstance(outputs, tuple):
return tuple(np_ify(x) for x in outputs)
else:
return np_ify(outputs)
def torch_ify(np_array_or_other):
if isinstance(np_array_or_other, np.ndarray):
return from_numpy(np_array_or_other)
else:
return np_array_or_other
def np_ify(tensor_or_other):
if isinstance(tensor_or_other, torch.autograd.Variable):
return get_numpy(tensor_or_other)
else:
return tensor_or_other
def _elem_or_tuple_to_variable(elem_or_tuple):
if isinstance(elem_or_tuple, tuple):
return tuple(
_elem_or_tuple_to_variable(e) for e in elem_or_tuple
)
return from_numpy(elem_or_tuple).float()
def _filter_batch(np_batch):
for k, v in np_batch.items():
if v.dtype == np.bool:
yield k, v.astype(int)
else:
yield k, v
def np_to_pytorch_batch(np_batch):
return {
k: _elem_or_tuple_to_variable(x)
for k, x in _filter_batch(np_batch)
if x.dtype != np.dtype('O') # ignore object (e.g. dictionaries)
}