Problem using self-defined convolutional network on the heisenbergJ1J2 model

[sweetpotat5]

Hello there,

I'm trying to run the heisenbergJ1J2 example with my self-defined convolutional network. Here is my code:

# import basic libraries
import netket as nk
import numpy as np
import json
from math import pi

from netket.graph import Triangular
from netket.models import GCNN
from netket.utils.group.planar import C
from netket import experimental as nkx
import flax
import flax.linen as nn

L = 10

lattice = nk.graph.Square(L, max_neighbor_order=2)
hi = nk.hilbert.Spin(s=1 / 2, total_sz=0, N=lattice.n_nodes)
H = nk.operator.Heisenberg(hilbert=hi, graph=lattice, J=[1.0, 0.5])

# Find an approximate ground state
sampler = nk.sampler.MetropolisExchange(
    hilbert=hi,
    n_chains=1024,
    graph=lattice,
    d_max=2,
)

### use self-defined convolutional network
class CNN(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=2,kernel_size=(3,3))(x)
        x = nn.relu(x)
        x = jnp.sum(x, axis=-1)
        return x
machine = CNN()
###

opt = nk.optimizer.Sgd(learning_rate=0.02)
sr = nk.optimizer.SR(diag_shift=0.01)


vstate = nk.vqs.MCState(
    sampler=sampler,
    model=machine,
    n_samples=1024,
    n_discard_per_chain=0,
    chunk_size=4096,
)
gs = nk.driver.VMC(H, opt, variational_state=vstate, preconditioner=sr)
gs.run(n_iter=200, out="ground_state")

data = json.load(open("ground_state.log"))
print(np.mean(data["Energy"]["Mean"]["real"][-20:]) / 400)
print(np.std(data["Energy"]["Mean"]["real"][-20:]) / 400)

The only difference between my code and the example is that I used a self-defined convolutional network with only one layer. However, I'm getting an error message:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 43
     39 opt = nk.optimizer.Sgd(learning_rate=0.02)
     40 sr = nk.optimizer.SR(diag_shift=0.01)
---> 43 vstate = nk.vqs.MCState(
     44     sampler=sampler,
     45     model=machine,
     46     n_samples=1024,
     47     n_discard_per_chain=0,
     48     chunk_size=4096,
     49 )
     50 gs = nk.driver.VMC(H, opt, variational_state=vstate, preconditioner=sr)
     51 gs.run(n_iter=200, out="ground_state")

~/.local/lib/python3.9/site-packages/netket/vqs/mc/mc_state/state.py:256, in MCState.__init__(self, sampler, model, n_samples, n_samples_per_rank, n_discard, n_discard_per_chain, chunk_size, variables, init_fun, apply_fun, sample_fun, seed, sampler_seed, mutable, training_kwargs)
    254     self.variables = variables
    255 else:
--> 256     self.init(seed, dtype=sampler.dtype)
    258 if sampler_seed is None and seed is not None:
    259     key, key2 = jax.random.split(nkjax.PRNGKey(seed), 2)

~/.local/lib/python3.9/site-packages/netket/vqs/mc/mc_state/state.py:291, in MCState.init(self, seed, dtype)
    287 key = nkjax.PRNGKey(seed)
    289 dummy_input = jnp.zeros((1, self.hilbert.size), dtype=dtype)
--> 291 variables = jit_evaluate(self._init_fun, {"params": key}, dummy_input)
    292 self.variables = variables

    [... skipping hidden 14 frame]

~/.local/lib/python3.9/site-packages/netket/vqs/mc/mc_state/state.py:95, in jit_evaluate(fun, *args)
     86 @partial(jax.jit, static_argnums=0)
     87 def jit_evaluate(fun: Callable, *args):
     88     """
     89     call `fun(*args)` inside of a `jax.jit` frame.
     90 
   (...)
     93         args: the arguments to the function.
     94     """
---> 95     return fun(*args)

~/.local/lib/python3.9/site-packages/netket/vqs/mc/mc_state/state.py:195, in MCState.__init__.<locals>.<lambda>(model, *args, **kwargs)
    190     _, model = maybe_wrap_module(model)
    192     self._model = model
    194     self._init_fun = nkjax.HashablePartial(
--> 195         lambda model, *args, **kwargs: model.init(*args, **kwargs), model
    196     )
    197     self._apply_fun = wrap_to_support_scalar(
    198         nkjax.HashablePartial(
    199             lambda model, pars, x, **kwargs: model.apply(pars, x, **kwargs),
    200             model,
    201         )
    202     )
    204 elif apply_fun is not None:

    [... skipping hidden 9 frame]

Cell In[1], line 32, in CNN.__call__(self, x)
     30 @nn.compact
     31 def __call__(self, x):
---> 32     x = nn.Conv(features=2,kernel_size=(3,3))(x)
     33     x = nn.relu(x)
     34     x = jnp.sum(x, axis=-1)

    [... skipping hidden 2 frame]

~/.local/lib/python3.9/site-packages/flax/linen/linear.py:438, in _Conv.__call__(self, inputs)
    436 inputs, kernel, bias = promote_dtype(inputs, kernel, bias, dtype=self.dtype)
    437 if self.shared_weights:
--> 438   y = lax.conv_general_dilated(
    439       inputs,
    440       kernel,
    441       strides,
    442       padding_lax,
    443       lhs_dilation=input_dilation,
    444       rhs_dilation=kernel_dilation,
    445       dimension_numbers=dimension_numbers,
    446       feature_group_count=self.feature_group_count,
    447       precision=self.precision
    448   )
    449 else:
    450   y = lax.conv_general_dilated_local(
    451       lhs=inputs,
    452       rhs=kernel,
   (...)
    459       precision=self.precision
    460   )

    [... skipping hidden 2 frame]

~/.local/lib/python3.9/site-packages/jax/_src/lax/lax.py:4513, in _ceil_divide(x1, x2)
   4512 def _ceil_divide(x1, x2):
-> 4513   return -np.floor_divide(np.negative(x1), x2)

ValueError: operands could not be broadcast together with shapes (0,) (2,) 

I'm having trouble understanding where the issue could come from. I'm using netket v3.6. Any comments would be appreciated.

Thank you!

[jobdky]

Hi @tianshuhuang98,

I am not completely sure, but I think you get an error because nn.Conv expects a feature dimension of the input.
This means that in your CNN you should first expand the shape of the input, so the last dimension corresponds to the number of features. You should also keep in mind that the inputs have the shape (batch_size, lattice.n_nodes) so far and not (batch_size, L, L). To use the 2D convolution I think you have to reshape the inputs as x = x.reshape(x.shape[0], L, L, 1) before applying the convolution.
The summation for the output should then include a sum over the last three dimensions of x.

Furthermore, you should use complex weights in your network for a proper approximation of the ground state (which would require a deeper network). This can be done by specifying x = nn.Conv(features=2,kernel_size=(3,3), param_dtype = jnp.complex128)(x).

[sweetpotat5]

This solves the problem, thank you!