Hi @riccardomurri,
I think I have good news 
. I was able to optimize your 2D convolution quite a bit, and now it seems pretty fast. It’s a bit unfortunate that while in Julia generally, naive loops are fast, in Convex.jl things are implemented in a vectorized way and sometimes one needs to employ tricks to get good performance. Now with N=10, m=10, r=2, the code runs in 15 seconds on my laptop, whereas before it did not finish in 10 minutes. But of course, we aren’t even solving anything yet, so maybe the performance will still not be good enough. I also haven’t improved the hashing time yet; I’m actually a bit more interested in the long-term solution there of not needing special hashing in the first place, though if there is a good short term fix, I’d take that too!
Step 1: Write a 2D convolution in terms of a 1D one
We will start with a numerical expression that doesn’t yet work with Convex.jl variables.
# From `test/definitions.jl` in Convex.jl
# a 1D convolution
function conv1D(h, x)
m = length(h)
n = length(x)
zero_pad_x(i) = 1 <= i <= n ? x[i] : 0
[ sum(h[j]*zero_pad_x(i-j+1) for j = 1:m) for i = 1:m+n-1 ]
end
# Perform a 2D convolution in terms of a 1D one
# Following https://sites.ualberta.ca/~mostafan/Files/Papers/md_convolution_TLE2009.pdf
function conv2D(X::AbstractMatrix, Y::AbstractMatrix)
M, N = size(X)
K, L = size(Y)
Z_rows = M + K - 1
Z_cols = N + L -1
X_padded = zeros(Z_rows, Z_cols)
X_padded[1:size(X, 1), 1:size(X,2)] .= X
X_vector = vec(X_padded)[1: (N-1)*(M+K-1) + M]
Y_padded = zeros(Z_rows, Z_cols)
Y_padded[1:size(Y, 1), 1:size(Y,2)] .= Y
Y_vector = vec(Y_padded)[1: ( L-1)*(M+K-1) + K]
Z_vector = conv1D(X_vector, Y_vector)
Z = reshape(Z_vector, Z_rows, Z_cols)
Z[M:K, N:L]
end
Now, let’s check this works against your conv2d:
function _conv2D(Q, x)
rows_Q, cols_Q = size(Q);
rows_x, cols_x = size(x);
i_min = rows_Q + 1; # such that i-h>=1 when i=i_min and h=rows_Q
j_min = cols_Q + 1; # such that j-k>=1 when j=j_min and k=cols_Q
i_max = rows_x + 1; # such that i-h<=rows_x when i=i_max and h=1
# this was corrected from `j_max = rows_x + 1`
j_max = cols_x + 1;
#
return [
reduce(+, (Q[h,k]*x[i-h, j-k]
for h in 1:rows_Q, k in 1:cols_Q))
for i in i_min:i_max, j in j_min:j_max
];
end
using Test
@testset "conv2D implementation" begin
for (M,N,K,L) in [ rand(2:100, 4) for _ = 1:5]
Q = rand(M, N)
x = rand(K, L)
@test _conv2D(Q, x) ≈ conv2D(Q,x)
end
end
Test Summary: | Pass Total
conv2D implementation | 5 5
Step 2: modify our 2D convolution so it can be used with Convex.jl variables
using Convex
"""
Convex_conv2D(X::AbstractMatrix, Y::Convex.AbstractExpr) -> Z, constraints
Produces the 2D convolution `Z` of `X` and `Y`, along with a set
of constraints that must be added to the problem.
"""
function Convex_conv2D(X::AbstractMatrix, Y::Convex.AbstractExpr)
M, N = size(X)
K, L = size(Y)
Z_rows = M + K - 1
Z_cols = N + L -1
X_padded = zeros(Z_rows, Z_cols)
X_padded[1:size(X, 1), 1:size(X,2)] .= X
X_vector = vec(X_padded)[1: (N-1)*(M+K-1) + M]
Y_padded = Variable(Z_rows, Z_cols)
constraints = [
Y_padded[1:size(Y, 1), 1:size(Y,2)] == Y,
Y_padded[size(Y,1)+1:end, :] == 0,
Y_padded[1:size(Y,1), size(Y,2)+1:end] == 0,
]
Y_vector = vec(Y_padded)[1: ( L-1)*(M+K-1) + K]
Z_vector = Convex.conv(X_vector, Y_vector)
Z = reshape(Z_vector, Z_rows, Z_cols)
Z[M:K, N:L], constraints
end
Note our 2D convolution produces the expression, but also a vector of constraints we need to add to the problem later on.
Step 3: Optimization
We notice that the 1D convolution in Convex.jl’s source creates a sparse matrix based on the first argument of the convolution. Since this doesn’t change over the course of the loop, we can just create it once and reuse it. We can also create it much faster than is done in Convex.jl’s source.
using SparseArrays
# The following is modified from Convex.jl's source
"""
conv1D_matrix(v::AbstractVector, n::Int) -> SparseMatrixCSC
Create a sparse matrix `A` such that if `x` has length `n`,
then we have `A * x ≈ conv1D(v, x)`.
"""
function conv1D_matrix(v::AbstractVector, n::Int)
m = length(v)
# It is much more efficient to construct sparse matrices
# this way rather than starting from `spzeros` and indexing into it.
# I should make a PR to Convex.jl to fix this there.
I = Int[]
J = Int[]
V = eltype(v)[]
for i = 1:n
for (ind,j) = enumerate(i:m+i-1: i)
push!(I, i)
push!(J, j)
push!(V, v[ind])
end
end
A = sparse(I,J,V, m + n - 1, n)
return A
end
# A helper to create the correct matrix that performs
# the 1D convolution
function make_A(Q, sz)
M, N = size(Q)
K, L = sz
Z_rows = M+K-1
Z_cols = N + L -1
Qpad = zeros(Z_rows, Z_cols)
Qpad[1:size(Q, 1), 1:size(Q,2)] .= Q
Qv = vec(Qpad)[1: (N-1)*(M+K-1) + M]
conv1D_matrix(Qv, (L-1)*(M+K-1) + K)
end
# Use the `A` matrix. Note we only depend on the
# size of the first argument now, not the value,
# since that is encoded in the `A` matrix.
function Convex_conv2D_matrix(X_size::Tuple, Y::Convex.AbstractExpr, A::AbstractMatrix)
M, N = X_size
K, L = size(Y)
Z_rows = M + K - 1
Z_cols = N + L -1
Y_padded = Variable(Z_rows, Z_cols)
constraints = [
Y_padded[1:size(Y, 1), 1:size(Y,2)] == Y,
Y_padded[size(Y,1)+1:end, :] == 0,
Y_padded[1:size(Y,1), size(Y,2)+1:end] == 0,
]
Y_vector = vec(Y_padded)[1: ( L-1)*(M+K-1) + K]
Z_vector = A * Y_vector
Z = reshape(Z_vector, Z_rows, Z_cols)
Z[M:K, N:L], constraints
end
Step 4: put it all together
In the following, I’ve kept my TimerOutputs logging which I used originally to see that the convolution was the bottleneck.
using SparseArrays: sprand;
using Dates: now;
using TimerOutputs
function make_constraints(N, r, m)
to = TimerOutput()
s = 4*r + 1;
x = [Variable(m*s+2r, m*s+2r) for _ in 1:N];
y = zeros(s,s,N);
# for demoing, we can assume `y` is constant in time
for t = 1:N
y[:, :, t] .= sprand(Float64, s, s, 0.15);
end
@timeit to "make Q" begin
Q = zeros(2r+1,2r+1);
for i=1:(2r+1), j=1:(2r+1)
d = (i-r-1)^2 + (j-r-1)^2
f = r/3;
Q[i,j] = exp(-4*log(2)*d/f^2)
end
end
@timeit to "make A" A = make_A(Q, (m*s+2r, m*s+2r))
± = [-1.0, +1.0]; # well, Julia does allow crazy variable
totals = Array{Any}(undef, N);
constraints = Constraint[]
extra_constraints = Constraint[]
for t in 1:N
@info "Starting computation of constraint A[$(t)] at $(now()) ..."
# 1. introduce sign perturbation `R`
@timeit to "a1" a1 = (rand(±, size(x[t])) .* x[t]) + x[1];
@timeit to "a2" a2, cs = Convex_conv2D_matrix(size(Q), a1, A);
append!(constraints, cs)
# 3. project into lo-res space
@timeit to "a3" a3 = [
sum( (a2[I,J]
for I in (1+m*(i-1)):m*i,
J in (1+m*(j-1)):m*j))
for i in 1:s, j in 1:s
];
# fetch the corresponding lo-res vector out of `y`
@timeit to "a4" a4 = a3 - y[:, :, t];
# now add into the effective constraint
@timeit to "totals" totals[t] = sum((a4[i,j] * a4[i,j] for i in 1:s, j in 1:s));
end; # for t
# build and return final constraint
@timeit to "total" total = sum(totals);
print_timer(to)
return total, constraints
end
Note that this produces total, like the original code, but also a vector of constraints that must be added to the problem.
Step 5: compare vs original
Lastly, we can compare versus the original version
function make_constraints_original(N, r, m)
s = 4*r + 1;
x = [Variable(m*s+2r, m*s+2r) for _ in 1:N];
y = zeros(s,s,N);
# for demoing, we can assume `y` is constant in time
for t = 1:N
y[:, :, t] .= sprand(Float64, s, s, 0.15);
end
Q = zeros(2r+1,2r+1);
for i=1:(2r+1), j=1:(2r+1)
d = (i-r-1)^2 + (j-r-1)^2
f = r/3;
Q[i,j] = exp(-4*log(2)*d/f^2)
end
function conv2d_original(Q, x)
rows_Q, cols_Q = size(Q);
rows_x, cols_x = size(x);
i_min = rows_Q + 1; # such that i-h>=1 when i=i_min and h=rows_Q
j_min = cols_Q + 1; # such that j-k>=1 when j=j_min and k=cols_Q
i_max = rows_x + 1; # such that i-h<=rows_x when i=i_max and h=1
j_max = rows_x + 1; # such that j-k<=rows_x when j=j_max and k=1
return [
reduce(+, (Q[h,k]*x[i-h, j-k]
for h in 1:rows_Q, k in 1:cols_Q))
for i in i_min:i_max, j in j_min:j_max
];
end
± = [-1.0, +1.0]; # well, Julia does allow crazy variable
totals = Array{Any}(undef, N);
for t in 1:N
@info "Starting computation of constraint A[$(t)] at $(now()) ..."
# 1. introduce sign perturbation `R`
a1 = (rand(±, size(x[t])) .* x[t]) + x[1];
# 2. 2D-convolution with PSF matrix `Q`
a2 = conv2d_original(Q, a1);
# 3. project into lo-res space
a3 = [
reduce(+, (a2[I,J]
for I in (1+m*(i-1)):m*i,
J in (1+m*(j-1)):m*j))
for i in 1:s, j in 1:s
];
# fetch the corresponding lo-res vector out of `y`
a4 = a3 - y[:, :, t];
# now add into the effective constraint
totals[t] = reduce(+, (a4[i,j] * a4[i,j] for i in 1:s, j in 1:s));
end; # for t
# build and return final constraint
total = reduce(+, totals);
end
Then
N = 2
r = 5
m = 2
@time make_constraints_original(N, r, m);
@time make_constraints(N, r, m);
gives me 17 seconds vs 2.5. More importantly, setting N=r=10, the second completes in 15 seconds.
Now, I hope I haven’t made any errors during this optimization. I tested my numerical 2D convolution, as shown above, but I don’t have a full optimization problem to compare the two more fully. So please test it (on small problems) against your original code and let me know if it works.
If it does work, it might be nice to add something like this to Convex.jl. The bit where there’s extra constraints to carry around makes for a slightly awkward API, but luckily `AbstractVariable` interface; allow variables to carry constraints by ericphanson · Pull Request #358 · jump-dev/Convex.jl · GitHub will solve that by allowing us to tell Z to carry its own constraints. This actually makes for some nice motivation for that feature.
wouldn’t 
Given what