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
|
import torch
class TorchTensorEngine(object):
def rand(self, shape, device=None, dtype=None, requires_grad=False):
return torch.rand(shape, device=device, dtype=dtype, requires_grad=requires_grad)
def randn(self, shape, device=None, dtype=None, requires_grad=False):
return torch.randn(shape, device=device, dtype=dtype, requires_grad=requires_grad)
def nchw_rand(self, shape, device=None, requires_grad=False):
return self.rand(shape, device=device, requires_grad=requires_grad)
def reset(self, _):
pass
def rand_like(self, v):
return torch.rand_like(v)
def numpy(self, t):
return t.cpu().numpy()
def mul(self, t1, t2):
return t1 * t2
def add(self, t1, t2):
return t1 + t2
def batch_norm(self, data, mean, var, training):
return torch.nn.functional.batch_norm(data, mean, var, training=training)
def instance_norm(self, data):
return torch.nn.functional.instance_norm(data)
def layer_norm(self, data, shape):
return torch.nn.functional.layer_norm(data, shape)
def sync_cuda(self):
torch.cuda.synchronize()
def backward(self, tensors, grad_tensors, _):
torch.autograd.backward(tensors, grad_tensors=grad_tensors)
def sum(self, data, dims):
return torch.sum(data, dims)
def softmax(self, data, dim=None, dtype=None):
return torch.nn.functional.softmax(data, dim, dtype)
def cat(self, inputs, dim=0):
return torch.cat(inputs, dim=dim)
def clamp(self, data, min, max):
return torch.clamp(data, min=min, max=max)
def relu(self, data):
return torch.nn.functional.relu(data)
def tanh(self, data):
return torch.tanh(data)
def max_pool2d(self, data, kernel_size, stride=1):
return torch.nn.functional.max_pool2d(data, kernel_size, stride=stride)
def avg_pool2d(self, data, kernel_size, stride=1):
return torch.nn.functional.avg_pool2d(data, kernel_size, stride=stride)
def conv2d_layer(self, ic, oc, kernel_size, groups=1):
return torch.nn.Conv2d(ic, oc, kernel_size, groups=groups)
def matmul(self, t1, t2):
return torch.matmul(t1, t2)
def to_device(self, module, device):
return module.to(device)
|
