Closed liamgd closed 1 year ago
First try:
def prob_create(df: int, device: torch.device, dtype: torch.dtype):
@torch.jit.script
def prob(value: torch.Tensor, dft, scalar, exponent):
return (value ** 2.0 / dft + 1) ** exponent / scalar ** exponent
dft = torch.tensor(df, device=device, dtype=dtype)
scalar = torch.tensor(
math.sqrt(df)
* math.sqrt(math.pi)
* math.exp(math.lgamma(0.5 * df) - math.lgamma(0.5 * df + 0.5)),
device=device,
dtype=dtype,
)
exponent = torch.tensor(-0.5 * (df + 1.0), device=device, dtype=dtype)
return lambda value: prob(value, dft, scalar, exponent)
Computation overflows and results in a tensor of infinities.
One equivalent algorithm:
@torch.jit.script
class Prob:
def __init__(
self, df: int, device: torch.device, dtype: torch.dtype
) -> None:
self.df = df
self.offset = torch.tensor(
-0.5 * math.log(df)
- 0.5 * math.log(math.pi)
- math.lgamma(0.5 * df)
+ math.lgamma(0.5 * (df + 1.0)),
device=device,
dtype=dtype,
)
self.scalar = torch.tensor(
0.5 * (self.df + 1.0), device=device, dtype=dtype
)
def prob(self, value: torch.Tensor):
return (
self.offset - torch.log1p(value ** 2.0 / self.df) * self.scalar
).exp()
For torch just in time compilation using the @torch.jit.script
decorator, a class must be used to accept a variable df.
(because torch JIT does not support running with line_profiler)
Comparing a torch.distributions.studentT.StudentT(df).log_prob(value).exp()
with the new Prob algorithm over 100 repeats for accuracy. Test data: torch.random 100,000,000 by 4 tensor.
Total time: 13.1775 s
File: .\tecpg\regression_full.py
Function: test_prob at line 259
Line # Hits Time Per Hit % Time Line Contents
==============================================================
259 @profile
260 def test_prob() -> None:
261 1 6.9 6.9 0.0 torch.cuda.empty_cache()
262 1 0.2 0.2 0.0 df = 200
263 1 12.2 12.2 0.0 device = torch.device('cuda')
264 1 0.7 0.7 0.0 dtype = torch.float32
265 1 0.5 0.5 0.0 prob_one = lambda t: torch.distributions.StudentT(df).log_prob(t).exp()
266 1 63757.1 63757.1 0.5 prob_two = Prob(df, device, dtype).prob
267 1 848254.1 848254.1 6.4 test = torch.rand((100_000_000, 4), device=device, dtype=dtype)
268 1 311.1 311.1 0.0 print(torch.cuda.memory_allocated())
269 100 63.3 0.6 0.0 for _ in range(100):
270 100 12254983.8 122549.8 93.0 prob_one(test)
271 100 10153.0 101.5 0.1 prob_two(test)
With torch.distributions.studentT.StudentT(df).log_prob(value).exp()
: 122549.8 μs average per loop
With the new Prob algorithm, not JIT: 101.5 average per loop
The new algorithm is approximately 1207.39 times faster from this test.
Tests of combinations of region filtration, p-value filtration, and chunking with and without @torch.jit.script on the probability function. 300 samples, 50,000 methylation loci, and 10,000 gene expression loci on a 2070 super with 8 GB VRAM.
The first number after each command is the time without torch JIT, and the second is with torch JIT.
tecpg run mlr-full: *Too much GPU memory allocation*
tecpg run mlr-full --cis: 93.1929 to 98.289 seconds
tecpg run mlr-full --distal: 94.3363 to 98.0763 seconds
tecpg run mlr-full --trans: *Too much GPU memory allocation*
tecpg run mlr-full -p 0.05: 93.2036 to 95.8579 seconds
tecpg run mlr-full -p 0.05 --cis: 93.427 to 97.4576 seconds
tecpg run mlr-full -p 0.05 --distal: 93.5236 to 97.2995 seconds
tecpg run mlr-full -p 0.05 --trans: 93.5954 to 96.8116 seconds
tecpg run mlr-full -l 2000: 169.2086 to 169.5929 seconds
tecpg run mlr-full -l 2000 --cis: 102.665 to 101.5853 seconds
tecpg run mlr-full -l 2000 --distal: 104.6743 to 107.7221 seconds
tecpg run mlr-full -l 2000 --trans: 180.5167 to 175.9578 seconds
tecpg run mlr-full -l 2000 -p 0.05: 101.5071 to 102.3017 seconds
tecpg run mlr-full -l 2000 -p 0.05 --cis: 97.5001 to 97.8865 seconds
tecpg run mlr-full -l 2000 -p 0.05 --distal: 98.7852 to 98.6476 seconds
tecpg run mlr-full -l 2000 -p 0.05 --trans: 101.6967 to 100.2482 seconds
It seems that in general, for this input size, torch JIT does not have a significant positive impact on performance and instead offers a slight decrease in speed.
Tests with 300 samples, 5,000 methylation loci, and 5,000 gene expression loci.
The closure approach is incompatible with torch JIT, but from previous tests, JIT seems to slightly decrease performance. From the above tests, it is apparent that the closure solution is around 4% faster. Implemented in 9dae7ff.
Currently, the regression algorithms use
torch.exp
oftorch.distributions.studentT.StudentT(df).log_prob
. This could be improved.