<jupyter_start><jupyter_text>Define network<jupyter_code>class STM(nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        Initialize the attributes of STM.
        """
        super(STM, self).__init__() # inherit the __init__ from nn.Module.
        
        self.input=torch.rand(D_in)            
        self.linear1=nn.Linear(D_in, H)
        self.linear2=nn.Linear(H, D_out)
        self.x = self.init_x()            # initialize self.x with method init_x().
        
    def init_x(self):
        """
        Initialize self.x with two neurons.
        """
        return nn.Parameter(self.input)
        
    def forward(self):
        """
        Forward propagation.
        """
        h_relu = self.linear1(self.x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        
        return y_pred<jupyter_output><empty_output><jupyter_text>Predefined parameters<jupyter_code># Average energy
E_mean=10
# D_in=Input dimensions, H=Hidden dimension, D_out=Number of classes
D_in, H, D_out = 50, 100, 64
# Learning rate
lr=0.001
# Number of auglag
num_auglag=10
# Number of optimization
num_ops=1500<jupyter_output><empty_output><jupyter_text>Predefined data<jupyter_code># Lagrange multiplier
l = 0.0
# Penalty factor
mu = 1.0
# Penalty update factor
sigma = 1.5
# Penalty safeguard
mu_max=10
# Convergence tolerance
ec = 1e-7
# Tolerance factor of constraint
tolerance=0.01

# Artifical energy levels
E=np.array([(i+1)**2 for i in range(D_out)])
E=torch.tensor(E, requires_grad=False, dtype=torch.float32)

# Target distribution
y=torch.ones(D_out)/D_out<jupyter_output><empty_output><jupyter_text>Predefined functions<jupyter_code>softmax=torch.nn.Softmax(dim=-1)<jupyter_output><empty_output><jupyter_text>Building model<jupyter_code>model = STM(D_in, H, D_out)

loss_fn = torch.nn.MSELoss(reduction='sum')<jupyter_output><empty_output><jupyter_text>Plot probability distribution before training<jupyter_code># Data
init_dist = softmax(model()).detach().numpy()

# Plot
fig, ax = plt.subplots(figsize=(8,6))
ax.plot(init_dist)
ax.set_xlabel('State index', size=16)
ax.set_ylabel('Probability', size=16)
ax.set_ylim(0.0,0.03)
ax.tick_params(axis='both', which='major', direction='in', labelsize=16)
plt.show()<jupyter_output><empty_output><jupyter_text>Average energy before training<jupyter_code>print(torch.mean(E*softmax(model())))<jupyter_output>tensor(22.2884, grad_fn=<MeanBackward0>)<jupyter_text>Training<jupyter_code># Initial constraint
C_best=(torch.mean(E*softmax(model()))-E_mean).detach_().numpy()
aC_best=np.abs(C_best)

for j in range(num_auglag):
    
    if j!=0: print(j, loss.item())
    for i in range(num_ops):
    
        Y=model()
        loss = loss_fn(y,softmax(Y)) + 0.5*mu*(torch.mean(E*softmax(Y))-E_mean)**2\
               + l*(torch.mean(E*softmax(Y))-E_mean)
    
        model.zero_grad()
        loss.backward()
    
        with torch.no_grad():
            for param in model.parameters():
                param -= lr * param.grad
    
    C_p=(torch.mean(E*softmax(Y))-E_mean).detach_().numpy()
    aC_p=np.abs(C_p)
    if aC_p <= tolerance*aC_best:
        if aC_p <= ec: break
        # Update Lagrange multiplier.
        l += mu*C_p
        aC_best=aC_p
    else:
        # This is to make sure mu_{k+1}=sigma*mu_k but mu always less than or equal to mu_max.
        mu=min(sigma*mu, mu_max)<jupyter_output>1 0.013537404127418995
2 0.012619392946362495
3 0.011907356791198254
4 0.011332646012306213
5 0.010854380205273628
6 0.010447684675455093
7 0.010096314363181591
8 0.00978918094187975
9 0.009518287144601345<jupyter_text>Plot probability distribution after training<jupyter_code># Data
final_dist = softmax(model()).detach().numpy()

# Plot
fig, ax=plt.subplots(figsize=(8,6))
ax.plot(E.detach().numpy(), final_dist)
ax.set_xlabel('Energy', size=16)
ax.set_ylabel('Probability', size=16)
ax.set_yscale('log')
ax.tick_params(axis='both', which='both', direction='in', labelsize=16)
plt.show()<jupyter_output><empty_output><jupyter_text>Average energy after training<jupyter_code>torch.mean(E*softmax(model()))<jupyter_output><empty_output><jupyter_text>Constraint term after traiing<jupyter_code>torch.mean(E*softmax(model()))-E_mean<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1363850873053327, 'fraction_numerical': 0.059933931099575274, 'mean_word_length': 4.239802224969098, 'pattern_counts': {'":': 0, '<': 46, '<?xml version=': 0, '>': 44, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '465502', 'license': 'MIT', 'n_tokens_mistral': 1775, 'n_tokens_neox': 1621, 'n_words': 263, 'path': 'penalty_methods/DNN_augmented.ipynb', 'repo_name': 'tiwarylab/ps-LSTM', 'size': '4238'}