pytorch
79 строк · 2.3 Кб
1import torch
2
3
4class TorchTensorEngine:
5def rand(self, shape, device=None, dtype=None, requires_grad=False):
6return torch.rand(
7shape, device=device, dtype=dtype, requires_grad=requires_grad
8)
9
10def randn(self, shape, device=None, dtype=None, requires_grad=False):
11return torch.randn(
12shape, device=device, dtype=dtype, requires_grad=requires_grad
13)
14
15def nchw_rand(self, shape, device=None, requires_grad=False):
16return self.rand(shape, device=device, requires_grad=requires_grad)
17
18def reset(self, _):
19pass
20
21def rand_like(self, v):
22return torch.rand_like(v)
23
24def numpy(self, t):
25return t.cpu().numpy()
26
27def mul(self, t1, t2):
28return t1 * t2
29
30def add(self, t1, t2):
31return t1 + t2
32
33def batch_norm(self, data, mean, var, training):
34return torch.nn.functional.batch_norm(data, mean, var, training=training)
35
36def instance_norm(self, data):
37return torch.nn.functional.instance_norm(data)
38
39def layer_norm(self, data, shape):
40return torch.nn.functional.layer_norm(data, shape)
41
42def sync_cuda(self):
43torch.cuda.synchronize()
44
45def backward(self, tensors, grad_tensors, _):
46torch.autograd.backward(tensors, grad_tensors=grad_tensors)
47
48def sum(self, data, dims):
49return torch.sum(data, dims)
50
51def softmax(self, data, dim=None, dtype=None):
52return torch.nn.functional.softmax(data, dim, dtype)
53
54def cat(self, inputs, dim=0):
55return torch.cat(inputs, dim=dim)
56
57def clamp(self, data, min, max):
58return torch.clamp(data, min=min, max=max)
59
60def relu(self, data):
61return torch.nn.functional.relu(data)
62
63def tanh(self, data):
64return torch.tanh(data)
65
66def max_pool2d(self, data, kernel_size, stride=1):
67return torch.nn.functional.max_pool2d(data, kernel_size, stride=stride)
68
69def avg_pool2d(self, data, kernel_size, stride=1):
70return torch.nn.functional.avg_pool2d(data, kernel_size, stride=stride)
71
72def conv2d_layer(self, ic, oc, kernel_size, groups=1):
73return torch.nn.Conv2d(ic, oc, kernel_size, groups=groups)
74
75def matmul(self, t1, t2):
76return torch.matmul(t1, t2)
77
78def to_device(self, module, device):
79return module.to(device)
80