【コード解説】アフィン変換による動作追跡点の変形方法・Torch編

# Standard modules.
import gc
import sys
import time
from functools import partial

# CV/ML.
import numpy as np

import torch
import torch.nn as nn

print(f"Python:{sys.version}")
print(f"Numpy:{np.__version__}")
print(f"Torch:{torch.__version__}")

def get_perf_str(val):
    token_si = ["", "m", "µ", "n", "p"]
    exp_si = [1, 1e3, 1e6, 1e9, 1e12]
    perf_str = f"{val:3g}s"
    si = ""
    sval = val
    for token, exp in zip(token_si, exp_si):
        if val * exp > 1.0:
            si = token
            sval = val * exp
            break
    perf_str = f"{sval:3g}{si}s"
    return perf_str

def print_perf_time(intervals, top_k=None):
    if top_k is not None:
        intervals = np.sort(intervals)[:top_k]
    min = intervals.min()
    max = intervals.max()
    mean = intervals.mean()
    std = intervals.std()

    smin = get_perf_str(min)
    smax = get_perf_str(max)
    mean = get_perf_str(mean)
    std = get_perf_str(std)
    if top_k:
        print(f"Top {top_k} summary: Max {smax}, Min {smin}, Mean +/- Std {mean} +/- {std}")
    else:
        print(f"Overall summary: Max {smax}, Min {smin}, Mean +/- Std {mean} +/- {std}")

class PerfMeasure():
    def __init__(self,
                 trials=100,
                 top_k=10):
        self.trials = trials
        self.top_k = top_k

    def __call__(self, func):
        gc.collect()
        gc.disable()
        intervals = []
        for _ in range(self.trials):
            start = time.perf_counter()
            func()
            end = time.perf_counter()
            intervals.append(end - start)
        intervals = np.array(intervals)
        print_perf_time(intervals)
        if self.top_k:
            print_perf_time(intervals, self.top_k)
        gc.enable()
        gc.collect()

TRIALS = 100
TOPK = 10
pmeasure = PerfMeasure(TRIALS, TOPK)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Target device is {DEVICE}.")
JIT_OPT = False

def get_affine_matrix_2d_torch(center,
                               trans,
                               scale,
                               rot,
                               skew,
                               dtype=torch.float32):
    device = center.device
    center_m = torch.tensor([[1.0, 0.0, float(-center[0])],
                             [0.0, 1.0, float(-center[1])],
                             [0.0, 0.0, 1.0]], dtype=dtype, device=device)
    scale_m = torch.tensor([[float(scale[0]), 0.0, 0.0],
                            [0.0, float(scale[1]), 0.0],
                            [0.0, 0.0, 1.0]], dtype=dtype, device=device)
    _cos = torch.cos(rot)
    _sin = torch.sin(rot)
    rot_m = torch.tensor([[float(_cos), float(-_sin), 0.0],
                          [float(_sin), float(_cos), 0.0],
                          [0.0, 0.0, 1.0]], dtype=dtype, device=device)
    _tan = torch.tan(skew)
    skew_m = torch.tensor([[1.0, float(_tan[0]), 0.0],
                           [float(_tan[1]), 1.0, 0.0],
                           [0.0, 0.0, 1.0]], dtype=dtype, device=device)
    move = center + trans
    trans_m = torch.tensor([[1.0, 0.0, float(move[0])],
                            [0.0, 1.0, float(move[1])],
                            [0.0, 0.0, 1.0]], dtype=dtype, device=device)
    # Make affine matrix.
    mat = torch.eye(3, 3, dtype=dtype, device=device)
    mat = torch.matmul(center_m, mat)
    mat = torch.matmul(scale_m, mat)
    mat = torch.matmul(rot_m, mat)
    mat = torch.matmul(skew_m, mat)
    mat = torch.matmul(trans_m, mat)
    return mat.to(dtype)

def apply_affine_torch(inputs, mat):
    xy = inputs[:, :, :2]
    ones =  torch.ones([xy.shape[0], xy.shape[1], 1], device=inputs.device)
    xy = torch.cat([xy, ones], dim=-1)
    xy = torch.einsum("...j,ij", xy, mat)
    inputs[:, :, :2] = xy[:, :, :-1]
    return inputs

# Load data.
trackfile = "./finger_far0_non_static.npy"
reffile = "./finger_far0_non_static_affine.npy"
trackdata = np.load(trackfile).astype(np.float32)
refdata = np.load(reffile).astype(np.float32)
print(trackdata.shape)

# Remove person axis.
trackdata = trackdata[0]
refdata = refdata[0]

 # Get affine matrix.
center = torch.tensor([638.0, 389.0]).to(DEVICE)
trans = torch.tensor([100.0, 0.0]).to(DEVICE)
scale = torch.tensor([2.0, 0.5]).to(DEVICE)
rot = torch.tensor(np.radians(15.0)).to(DEVICE)
skew = torch.tensor(np.radians([15.0, 15.0])).to(DEVICE)
dtype = torch.float32
print("Parameters")
print("Center:", center)
print("Trans:", trans)
print("Scale:", scale)
print("Rot:", rot)
print("Skew:", skew)

def perf_wrap_func(trackdata, center, trans, scale, rot, skew, dtype):
    mat = get_affine_matrix_2d_torch(center, trans, scale, rot, skew, dtype=dtype)
    newtrack = apply_affine_torch(trackdata, mat)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
mat = get_affine_matrix_2d_torch(center, trans, scale, rot, skew, dtype=dtype)
newtrack = apply_affine_torch(testtrack, mat)
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

# Evaluate difference.
diff = (np.round(newtrack.detach().cpu().numpy()) - np.round(refdata)).sum()
print(f"Sum of error:{diff}")

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

print("Time after second call.")
target_fn = partial(perf_wrap_func,
                    trackdata=testtrack,
                    center=center, trans=trans, scale=scale, rot=rot, skew=skew,
                    dtype=dtype)
pmeasure(target_fn)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
mat = get_affine_matrix_2d_torch(center, trans, scale, rot, skew, dtype=dtype)
newtrack = apply_affine_torch(testtrack[:-1], mat)
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

print("Time after second call.")
target_fn = partial(perf_wrap_func,
                    trackdata=testtrack[:-1],
                    center=center, trans=trans, scale=scale, rot=rot, skew=skew,
                    dtype=dtype)
pmeasure(target_fn)

@torch.jit.script
def get_affine_matrix_2d_torch_jit(center: torch.Tensor,
                                   trans: torch.Tensor,
                                   scale: torch.Tensor,
                                   rot: torch.Tensor,
                                   skew: torch.Tensor,
                                   dtype: torch.dtype = torch.float32):
    ...

@torch.jit.script
def apply_affine_torch_jit(inputs, mat):
    ...

def perf_wrap_func(trackdata, center, trans, scale, rot, skew, dtype):
    mat = get_affine_matrix_2d_torch_jit(center, trans, scale, rot, skew, dtype=dtype)
    newtrack = apply_affine_torch_jit(trackdata, mat)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.perf_counter()
    mat = get_affine_matrix_2d_torch_jit(center, trans, scale, rot, skew, dtype=dtype)
    newtrack = apply_affine_torch_jit(testtrack, mat)
    interval = time.perf_counter() - start
    print_perf_time(np.array([interval]))

    # Evaluate difference.
    diff = (np.round(newtrack.detach().cpu().numpy()) - np.round(refdata)).sum()
    print(f"Sum of error:{diff}")

    testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

    print("Time after second call.")
    target_fn = partial(perf_wrap_func,
                        trackdata=testtrack,
                        center=center, trans=trans, scale=scale, rot=rot, skew=skew,
                        dtype=dtype)
    pmeasure(target_fn)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.perf_counter()
    mat = get_affine_matrix_2d_torch_jit(center, trans, scale, rot, skew, dtype=dtype)
    newtrack = apply_affine_torch_jit(testtrack[:-1], mat)
    interval = time.perf_counter() - start
    print_perf_time(np.array([interval]))

    # Evaluate difference.
    diff = (np.round(newtrack.detach().cpu().numpy()) - np.round(refdata[:-1])).sum()
    print(f"Sum of error:{diff}")

    testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

    print("Time after second call.")
    target_fn = partial(perf_wrap_func,
                        trackdata=testtrack[:-1],
                        center=center, trans=trans, scale=scale, rot=rot, skew=skew,
                        dtype=dtype)
    pmeasure(target_fn)

class RandomAffineTransform2D_Torch():
    def __init__(self,
                 center_joints,
                 apply_ratio,
                 trans_range,
                 scale_range,
                 rot_range,
                 skew_range,
                 random_seed=None,
                 device="cpu",
                 dtype=torch.float32):

        self.center_joints = center_joints
        if isinstance(self.center_joints, int):
            self.center_joints = [self.center_joints]

        self.apply_ratio = apply_ratio
        self.trans_range = trans_range
        self.scale_range = scale_range
        self.rot_range = np.radians(rot_range).tolist()
        self.skew_range = np.radians(skew_range).tolist()
        self.dtype = dtype
        self.device = device
        self.rng = torch.Generator(device=device)
        if random_seed is not None:
            self.rng.manual_seed(random_seed)

    def __call__(self, inputs):
        if torch.rand(1, generator=self.rng, device=self.device) >= self.apply_ratio:
            return inputs

        temp = inputs[:, self.center_joints, :]
        temp = temp.reshape([inputs.shape[0], -1, inputs.shape[-1]])
        mask = temp.sum(dim=(1, 2)) != 0
        # Use x and y only.
        center = temp[mask].mean(dim=0).mean(dim=0)[:2]

        # Random value in [0, 1].
        trans = torch.rand(2, generator=self.rng, device=self.device)
        scale = torch.rand(2, generator=self.rng, device=self.device)
        rot = torch.rand(1, generator=self.rng, device=self.device)
        skew = torch.rand(2, generator=self.rng, device=self.device)
        # Scale to target range.
        trans = (self.trans_range[1] - self.trans_range[0]) * trans + self.trans_range[0]
        scale = (self.scale_range[1] - self.scale_range[0]) * scale + self.scale_range[0]
        rot = (self.rot_range[1] - self.rot_range[0]) * rot + self.rot_range[0]
        skew = (self.skew_range[1] - self.skew_range[0]) * skew + self.skew_range[0]

        # Calculate matrix.
        mat = get_affine_matrix_2d_torch_jit(center, trans, scale, rot, skew,
            dtype=self.dtype)

        # Apply transform.
        inputs = apply_affine_torch_jit(inputs, mat)
        return inputs

aug_fn = RandomAffineTransform2D_Torch(
    center_joints=[11, 12],
    apply_ratio=1.0,
    trans_range=[-100.0, 100.0],
    scale_range=[0.5, 2.0],
    rot_range=[-30.0, 30.0],
    skew_range=[-30.0, 30.0],
    device=DEVICE,
    dtype=dtype)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.perf_counter()
    temp = aug_fn(testtrack)
    interval = time.perf_counter() - start
    print_perf_time(np.array([interval]))

    testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)
    print("Time after second call.")
    target_fn = partial(aug_fn, inputs=testtrack)
    pmeasure(target_fn)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.perf_counter()
    temp = aug_fn(testtrack[:-1])
    interval = time.perf_counter() - start
    print_perf_time(np.array([interval]))

    testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)
    print("Time after second call.")
    target_fn = partial(aug_fn, inputs=testtrack[:-1])
    pmeasure(target_fn)

class RandomAffineTransform2D_TorchModule(nn.Module):
    def __init__(self,
                 center_joints,
                 apply_ratio,
                 trans_range,
                 scale_range,
                 rot_range,
                 skew_range,
                 random_seed=None,
                 device="cpu",
                 dtype=torch.float32):
        super().__init__()

        self.center_joints = center_joints
        if isinstance(self.center_joints, int):
            self.center_joints = [self.center_joints]

        self.apply_ratio = apply_ratio
        self.trans_range = trans_range
        self.scale_range = scale_range
        self.rot_range = np.radians(rot_range).tolist()
        self.skew_range = np.radians(skew_range).tolist()
        self.dtype = dtype
        self.device = device
        # self.rng = torch.Generator(device=device)
        # if random_seed is not None:
        #     self.rng.manual_seed(random_seed)
        self.rng = None

    def forward(self, inputs):
        if torch.rand(1, generator=self.rng, device=self.device) >= self.apply_ratio:
            return inputs

        temp = inputs[:, self.center_joints, :]
        temp = temp.reshape([inputs.shape[0], -1, inputs.shape[-1]])
        mask = temp.sum(dim=(1, 2)) != 0
        # Use x and y only.
        center = temp[mask].mean(dim=0).mean(dim=0)[:2]

        # Random value in [0, 1].
        trans = torch.rand(2, generator=self.rng, device=self.device)
        scale = torch.rand(2, generator=self.rng, device=self.device)
        rot = torch.rand(1, generator=self.rng, device=self.device)
        skew = torch.rand(2, generator=self.rng, device=self.device)
        # Scale to target range.
        trans = (self.trans_range[1] - self.trans_range[0]) * trans + self.trans_range[0]
        scale = (self.scale_range[1] - self.scale_range[0]) * scale + self.scale_range[0]
        rot = (self.rot_range[1] - self.rot_range[0]) * rot + self.rot_range[0]
        skew = (self.skew_range[1] - self.skew_range[0]) * skew + self.skew_range[0]

        # Calculate matrix.
        mat = get_affine_matrix_2d_torch_jit(center, trans, scale, rot, skew,
            dtype=self.dtype)

        # Apply transform.
        inputs = apply_affine_torch_jit(inputs, mat)
        return inputs

aug_fn = RandomAffineTransform2D_TorchModule(
    center_joints=[11, 12],
    apply_ratio=1.0,
    trans_range=[-100.0, 100.0],
    scale_range=[0.5, 2.0],
    rot_range=[-30.0, 30.0],
    skew_range=[-30.0, 30.0],
    device=DEVICE,
    dtype=dtype)
aug_fn = torch.jit.script(aug_fn)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.perf_counter()
    temp = aug_fn(testtrack)
    interval = time.perf_counter() - start
    print_perf_time(np.array([interval]))

    testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)
    print("Time after second call.")
    target_fn = partial(aug_fn, inputs=testtrack)
    pmeasure(target_fn)

testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)

with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.perf_counter()
    temp = aug_fn(testtrack[:-1])
    interval = time.perf_counter() - start
    print_perf_time(np.array([interval]))

    testtrack = torch.tensor(trackdata.copy().astype(np.float32)).to(DEVICE)
    print("Time after second call.")
    target_fn = partial(aug_fn, inputs=testtrack[:-1])
    pmeasure(target_fn)