Skip to content

Commit 4e4cecc

Browse files
authored
Merge pull request #51 from TensorBFS/jg/refactor-uai-reading
refactor uai reading
2 parents 79797ea + dae58ad commit 4e4cecc

22 files changed

+330
-417
lines changed

Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "TensorInference"
22
uuid = "c2297e78-99bd-40ad-871d-f50e56b81012"
33
authors = ["Jin-Guo Liu", "Martin Roa Villescas"]
4-
version = "0.1.0"
4+
version = "0.2.0"
55

66
[deps]
77
Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

benchmark/bench_map.jl

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,10 @@ using Artifacts
66

77
const SUITE = BenchmarkGroup()
88

9-
model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
10-
problem = read_instance(model_filepath; evidence_filepath, solution_filepath)
9+
problem = problem_from_artifact("uai2014", "MAR" "Promedus", 14)
1110

1211
optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
13-
tn = TensorNetworkModel(problem; optimizer)
12+
tn = TensorNetworkModel(read_model(problem); optimizer, evidence=get_evidence(problem))
1413
SUITE["map"] = @benchmarkable most_probable_config(tn)
1514

1615
end # module

benchmark/bench_mar.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ using Artifacts
88

99
const SUITE = BenchmarkGroup()
1010

11-
model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
12-
problem = read_instance(model_filepath; evidence_filepath, solution_filepath)
11+
model_filepath, evidence_filepath, _, solution_filepath = get_model_filepaths("Promedus_14", "MAR")
12+
problem = read_model(model_filepath; evidence_filepath, solution_filepath)
1313

1414
optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)
1515
tn1 = TensorNetworkModel(problem; optimizer)

benchmark/bench_mmap.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ using Artifacts
66

77
const SUITE = BenchmarkGroup()
88

9-
model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
10-
problem = read_instance(model_filepath; evidence_filepath, solution_filepath)
9+
model_filepath, evidence_filepath, _, solution_filepath = get_model_filepaths("Promedus_14", "MAR")
10+
problem = read_model(model_filepath; evidence_filepath, solution_filepath)
1111
optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
1212

1313
# Does not marginalize any var

docs/make.jl

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ using TensorInference: OMEinsum
33
using TensorInference.OMEinsum: OMEinsumContractionOrders
44
using Documenter, Literate
55

6-
# Literate
6+
# Literate Examples
77
const EXAMPLE_DIR = pkgdir(TensorInference, "examples")
88
const LITERATE_GENERATED_DIR = pkgdir(TensorInference, "docs", "src", "generated")
99
mkpath(LITERATE_GENERATED_DIR)
@@ -15,6 +15,15 @@ for each in readdir(EXAMPLE_DIR)
1515
Literate.markdown(input_file, workdir; execute=true)
1616
end
1717

18+
const EXTRA_JL = ["performance.jl"]
19+
const SRC_DIR = pkgdir(TensorInference, "docs", "src")
20+
for each in EXTRA_JL
21+
cp(joinpath(SRC_DIR, each), joinpath(LITERATE_GENERATED_DIR, each); force=true)
22+
input_file = joinpath(LITERATE_GENERATED_DIR, each)
23+
@info "building" input_file
24+
Literate.markdown(input_file, LITERATE_GENERATED_DIR; execute=true)
25+
end
26+
1827
DocMeta.setdocmeta!(TensorInference, :DocTestSetup, :(using TensorInference); recursive=true)
1928

2029
makedocs(;
@@ -36,7 +45,7 @@ makedocs(;
3645
"Asia network" => "generated/asia/main.md",
3746
],
3847
"UAI file formats" => "uai-file-formats.md",
39-
"Performance tips" => "performance.md",
48+
"Performance tips" => "generated/performance.md",
4049
"API" => [
4150
"Public" => "api/public.md",
4251
"Internal" => "api/internal.md"

docs/src/api/public.md

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ TreeSA
4141
MMAPModel
4242
RescaledArray
4343
TensorNetworkModel
44-
UAIInstance
44+
ArtifactProblemSpec
45+
UAIModel
4546
```
4647

4748
## Functions
@@ -55,13 +56,14 @@ marginals
5556
maximum_logp
5657
most_probable_config
5758
probability
58-
read_evidence_file
59-
read_instance
60-
read_instance_from_artifact
59+
dataset_from_artifact
60+
problem_from_artifact
61+
read_model
62+
read_evidence
63+
read_solution
64+
read_queryvars
6165
read_model_file
62-
read_solution_file
66+
read_evidence_file
6367
read_td_file
6468
sample
65-
set_evidence!
66-
set_query!
6769
```

docs/src/performance.jl

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# # Performance Tips
2+
# ## Optimize contraction orders
3+
4+
# Let us use a problem instance from the "Promedus" dataset of the UAI 2014 competition as an example.
5+
using TensorInference
6+
problem = problem_from_artifact("uai2014", "MAR", "Promedus", 11)
7+
model, evidence = read_model(problem), read_evidence(problem);
8+
9+
# Next, we select the tensor network contraction order optimizer.
10+
optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.3:100)
11+
12+
# Here, we choose the local search based [`TreeSA`](@ref) algorithm, which often finds the smallest time/space complexity and supports slicing.
13+
# One can type `?TreeSA` in a Julia REPL for more information about how to configure the hyper-parameters of the [`TreeSA`](@ref) method,
14+
# while the detailed algorithm explanation is in [arXiv: 2108.05665](https://arxiv.org/abs/2108.05665).
15+
# Alternative tensor network contraction order optimizers include
16+
# * [`GreedyMethod`](@ref) (default, fastest in searching speed but worst in contraction complexity)
17+
# * [`KaHyParBipartite`](@ref)
18+
# * [`SABipartite`](@ref)
19+
20+
tn = TensorNetworkModel(model; optimizer, evidence);
21+
22+
# The returned object `tn` contains a field `code` that specifies the tensor network with optimized contraction order. To check the contraction complexity, please type
23+
contraction_complexity(tn)
24+
25+
# The returned object contains log2 values of the number of multiplications, the number elements in the largest tensor during contraction and the number of read-write operations to tensor elements.
26+
27+
probability(tn)
28+
29+
# ## Using the slicing technique to reduce the memory cost
30+
31+
# For large scale applications, it is also possible to slice over certain degrees of freedom to reduce the space complexity, i.e.
32+
# loop and accumulate over certain degrees of freedom so that one can have a smaller tensor network inside the loop due to the removal of these degrees of freedom.
33+
# In the [`TreeSA`](@ref) optimizer, one can set `nslices` to a value larger than zero to turn on this feature.
34+
# As a comparison we slice over 5 degrees of freedom, which can reduce the space complexity by at most 5.
35+
# In this application, the slicing achieves the largest possible space complexity reduction 5, while the time and read-write complexity are only increased by less than 1,
36+
# i.e. the peak memory usage is reduced by a factor ``32``, while the (theoretical) computing time is increased by at a factor ``< 2``.
37+
optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.3:100, nslices=5)
38+
tn = TensorNetworkModel(model; optimizer, evidence);
39+
contraction_complexity(tn)
40+
41+
# ## Faster Tropical tensor contraction to speed up MAP and MMAP
42+
# No extra effort is required to enjoy the BLAS level speed provided by [`TropicalGEMM`](https://github.com/TensorBFS/TropicalGEMM.jl).
43+
# The benchmark in the `TropicalGEMM` repo shows this performance is close to the theoretical optimal value.
44+
# Its implementation on GPU is under development in Github repo [`CuTropicalGEMM.jl`](https://github.com/ArrogantGao/CuTropicalGEMM.jl) as a part of [Open Source Promotion Plan summer program](https://summer-ospp.ac.cn/).
45+
46+
# ## Working with GPUs
47+
# To upload the computation to GPU, you just add `using CUDA` before calling the `solve` function, and set the keyword argument `usecuda` to `true`.
48+
# ```julia
49+
# julia> using CUDA
50+
# [ Info: OMEinsum loaded the CUDA module successfully
51+
#
52+
# julia> marginals(tn; usecuda = true);
53+
# ```
54+
55+
# Functions support `usecuda` keyword argument includes
56+
# * [`probability`](@ref)
57+
# * [`log_probability`](@ref)
58+
# * [`marginals`](@ref)
59+
# * [`most_probable_config`](@ref)
60+
61+
# ## Benchmarks
62+
# Please check our [paper (link to be added)]().

docs/src/performance.md

Lines changed: 0 additions & 93 deletions
This file was deleted.

examples/asia/main.jl

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -51,12 +51,12 @@ using TensorInference
5151
# Load the ASIA network model from the `asia.uai` file located in the examples
5252
# directory. See [Model file format (.uai)](@ref) for a description of the
5353
# format of this file.
54-
instance = read_instance(pkgdir(TensorInference, "examples", "asia", "asia.uai"))
54+
model = read_model_file(pkgdir(TensorInference, "examples", "asia", "asia.uai"))
5555

5656
# ---
5757

5858
# Create a tensor network representation of the loaded model.
59-
tn = TensorNetworkModel(instance)
59+
tn = TensorNetworkModel(model)
6060

6161
# ---
6262

@@ -76,13 +76,9 @@ get_vars(tn)
7676
# ---
7777

7878
# Set an evidence: Assume that the "X-ray" result (variable 7) is positive.
79-
set_evidence!(instance, 7 => 0)
80-
81-
# ---
82-
8379
# Since setting an evidence may affect the contraction order of the tensor
8480
# network, recompute it.
85-
tn = TensorNetworkModel(instance)
81+
tn = TensorNetworkModel(model, evidence = Dict(7 => 0))
8682

8783
# ---
8884

@@ -107,8 +103,7 @@ logp, cfg = most_probable_config(tn)
107103
# Compute the most probable values of certain variables (e.g., 4 and 7) while
108104
# marginalizing over others. This is known as Maximum a Posteriori (MAP)
109105
# estimation.
110-
set_query!(instance, [4, 7])
111-
mmap = MMAPModel(instance)
106+
mmap = MMAPModel(model, evidence=Dict(7=>0), queryvars=[4,7])
112107

113108
# ---
114109

0 commit comments

Comments
 (0)