deepmd.pt.model.task.denoise

deepmd.pt.model.task.denoise#

Classes#

DenoiseNet

Base fitting provides the interfaces of fitting net.

Module Contents#

class deepmd.pt.model.task.denoise.DenoiseNet(feature_dim, ntypes, attn_head=8, prefactor=[0.5, 0.5], activation_function='gelu', **kwargs)[source]#

Bases: deepmd.pt.model.task.fitting.Fitting

Base fitting provides the interfaces of fitting net.

feature_dim[source]#

ntypes[source]#

attn_head[source]#

prefactor[source]#

lm_head[source]#

output_def()[source]#: Returns the output def of the fitting net.

forward(pair_weights, diff, nlist_mask, features, sw, masked_tokens: torch.Tensor | None = None)[source]#

Calculate the updated coord. Args: - coord: Input noisy coord with shape [nframes, nloc, 3]. - pair_weights: Input pair weights with shape [nframes, nloc, nnei, head]. - diff: Input pair relative coord list with shape [nframes, nloc, nnei, 3]. - nlist_mask: Input nlist mask with shape [nframes, nloc, nnei].

Returns:

denoised_coord: Denoised updated coord with shape [nframes, nloc, 3].