After you’ve updated the weights, add the following lines to your code:
for p in mdl.parameters():
p.data.clamp_(0)
Example:
import torch
import torch.nn as nn
x = torch.arange(10).view(-1,1)
y = -3 * x
class NN_Linear_Model(nn.Module):
def __init__(self):
super().__init__()
self.lm = nn.Linear(1,1)
def forward(self, X):
out = self.lm(X)
return out
mdl = NN_Linear_Model()
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(mdl.parameters(), lr=0.0001)
for i in range(100):
optimizer.zero_grad()
y_pred = mdl(x)
loss = loss_fn(y_pred, y)
loss.backward()
optimizer.step()
for p in mdl.parameters():
p.data.clamp_(0)
if i % 10 == 0:
print(f'loss is {loss} at iter {i}, weight: {list(mdl.parameters())[0].item()}')
list(mdl.parameters())
One caveat is your model may not converge to the optimal point as you are restricting where your parameters can go.