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
|
import torch
from models.rgb_part_net import RGBPartNet
P, K = 2, 4
N, T, C, H, W = P * K, 10, 3, 64, 32
def rand_x1_x2_y(n, t, c, h, w):
x1 = torch.rand(n, t, c, h, w)
x2 = torch.rand(n, t, c, h, w)
y = []
for p in range(P):
y += [p] * K
y = torch.as_tensor(y)
return x1, x2, y
def test_default_rgb_part_net_cuda():
rgb_pa = RGBPartNet()
rgb_pa = rgb_pa.cuda()
x1, x2, y = rand_x1_x2_y(N, T, C, H, W)
x1, x2, y = x1.cuda(), x2.cuda(), y.cuda()
rgb_pa.train()
loss, metrics = rgb_pa(x1, x2, y)
_, _, _, _ = metrics
assert loss.device == torch.device('cuda', torch.cuda.current_device())
assert tuple(loss.size()) == ()
assert isinstance(_, float)
rgb_pa.eval()
x = rgb_pa(x1, x2)
assert x.device == torch.device('cuda', torch.cuda.current_device())
assert tuple(x.size()) == (23, N, 256)
|
