Dear OCP team,
I read the ForceNet paper and noticed data rotation augmentation is used in the paper for getting better result. However, I did not see how is it applied in the base model provided as the transform is set to none. Is the transformation done in some other place? Thanks!
Hi!
ForceNet was originally developed outside the OCP repo but the model was later ported over upon submission. It was benchmarked in the current repo to ensure results are consistent, but like you mentioned may be missing some things.
Regarding the transformation, the original implementation is shared below - Random3DRotate:
class Random3DRotate(object):
def __init__(self, axis=None):
r"""
Rotate both pos and force
args:
axis (int): select which axis to rotate along. default: None (randomly rotate along all the axis)
"""
self.axis = axis
def __call__(self, data):
if self.axis is None:
# rotate along x-axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
matrix = [[1, 0, 0], [0, cos, sin], [0, -sin, cos]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
# rotate along y-axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
matrix = [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
# rotate along z-axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
matrix = [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
else:
axis = self.axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
if axis == 0:
matrix = [[1, 0, 0], [0, cos, sin], [0, -sin, cos]]
elif axis == 1:
matrix = [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]
else:
matrix = [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
return data
The transformation happens at the dataloader level, which by default sets transform=None. To modify this, you need to define the transform here to:
self.train_dataset = registry.get_dataset_class(
self.config["task"]["dataset"]
)(self.config["dataset"], Random3DRotate())
FWIW - The repo currently has a similar transformation. This is missing the force rotation as part of the call for it to be useable for training, but I can push that change soon.
Hope this helps! Let me know if you run into any issues.
Thanks for you reply! I have added the new transform class into transform.py and modified the self.train_dataset as you mentioned. However, the transformation did not seem to apply and remained as None in the train_dataset object. I have to set self.transform = Random3DRotate() (here)in order to set the transform in train_dataset object. Even though now the transform method is running but seems to be purely doing computation and not adding augmented data into epochs.
Thanks for your answering in advance!
I’m not able to reproduce your error. The scheme I proposed earlier is working fine on my end and is transforming as expected. To confirm, the changes here should look as follows:
Verifying the transform was set properly:
Feel free to share code snippets on your end so I can better assist you.
The transformation works perfectly now. It turns out that I forgot to change the corresponding script for the transformation. Thanks for your help!
