Training starting again in sampling code

Rajatavaa: <p><strong>Training code</strong><br>
``import os<br>
import torch<br>
from Prepro import Data_orig<br>
import torch.nn as nn<br>
import torch.nn.functional as F<br>
from sklearn.model_selection import train_test_split</p>
<p>device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
print(“Using device:”, device)</p>
<p>class RNNModel(nn.Module):<br>
def <strong>init</strong>(self, vocab_size, n_embd, n_hidden1, n_hidden2, block_size):<br>
super(RNNModel, self).<strong>init</strong>()<br>
self.embedding = nn.Embedding(vocab_size, n_embd)<br>
self.rnn = nn.RNN(input_size=n_embd, hidden_size=n_hidden1, num_layers=2, batch_first=True, nonlinearity=‘tanh’)<br>
self.fc1 = nn.Linear(n_hidden1 * block_size, n_hidden2)<br>
self.fc2 = nn.Linear(n_hidden2, vocab_size)<br>
self._init_weights()</p>
<pre><code>def _init_weights(self):
    nn.init.normal_(self.embedding.weight, mean=0, std=0.1)
    nn.init.xavier_uniform_(self.fc1.weight)
    nn.init.xavier_uniform_(self.fc2.weight)

def forward(self, x):
    x = self.embedding(x)  # (batch_size, block_size, n_embd)
    x, _ = self.rnn(x)  # (batch_size, block_size, n_hidden1)
    x = x.reshape(x.size(0), -1)  # (batch_size, n_hidden1 * block_size)
    x = F.tanh(self.fc1(x))  # (batch_size, n_hidden2)
    x = self.fc2(x)          # (batch_size, vocab_size)
    return x
</code></pre>
<h1><a name="p-459023-model-parameters-1" class="anchor" href="#p-459023-model-parameters-1"></a>Model parameters</h1>
<p>n_embd = 10<br>
n_hidden1 = 500<br>
n_hidden2 = 200<br>
block_size = 3<br>
g = torch.Generator().manual_seed(2147483647)</p>
<h1><a name="p-459023-load-the-dataset-2" class="anchor" href="#p-459023-load-the-dataset-2"></a>Load the dataset</h1>
<p>dataset = Data_orig()<br>
X_train, X_test, y_train, y_test = train_test_split(dataset.X, dataset.Y,test_size=0.1, random_state=42)</p>
<h1><a name="p-459023-initialize-the-model-3" class="anchor" href="#p-459023-initialize-the-model-3"></a>Initialize the model</h1>
<p>model = RNNModel(vocab_size=dataset.vocab_size, n_embd=n_embd, n_hidden1=n_hidden1, n_hidden2=n_hidden2, block_size=block_size).to(device)</p>
<h1><a name="p-459023-define-optimizer-loss-function-and-scheduler-4" class="anchor" href="#p-459023-define-optimizer-loss-function-and-scheduler-4"></a>Define optimizer, loss function, and scheduler</h1>
<p>optimizer = torch.optim.Adam(model.parameters(), lr=0.001)<br>
criterion = nn.CrossEntropyLoss()<br>
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10000, gamma=0.5)</p>
<h1><a name="p-459023-create-a-directory-to-save-checkpoints-5" class="anchor" href="#p-459023-create-a-directory-to-save-checkpoints-5"></a>Create a directory to save checkpoints</h1>
<p>checkpoint_dir = “checkpoints”<br>
os.makedirs(checkpoint_dir, exist_ok=True)</p>
<h1><a name="p-459023-training-loop-with-checkpoint-saving-6" class="anchor" href="#p-459023-training-loop-with-checkpoint-saving-6"></a>Training loop with checkpoint saving</h1>
<p>max_steps = 100000<br>
batch_size = 50<br>
checkpoint_interval = 10000  # Save every 10,000 steps</p>
<p>for i in range(max_steps):<br>
# Mini-batch sampling<br>
ix = torch.randint(0, X_train.shape[0], (batch_size,))<br>
x_batch = X_train[ix].to(device)<br>
y_batch = y_train[ix].to(device)</p>
<pre><code># Forward pass
optimizer.zero_grad()
outputs = model(x_batch)
loss = criterion(outputs, y_batch)

# Backward pass and optimization
loss.backward()
optimizer.step()
scheduler.step()

# Track stats and print progress
if i % checkpoint_interval == 0:
    print(f'Step {i}/{max_steps}: Loss = {loss.item():.4f}, Learning Rate = {scheduler.get_last_lr()[0]:.6f}')

# Save checkpoint
if i % checkpoint_interval == 0 and i &gt; 0:
    checkpoint_path = os.path.join(checkpoint_dir, f'model_step_{i}.pth')
    torch.save({
        'step': i,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'loss': loss.item(),
    }, checkpoint_path)
    print(f"Checkpoint saved at step {i} to {checkpoint_path}")
</code></pre>
<h1><a name="p-459023-final-save-at-the-end-of-training-7" class="anchor" href="#p-459023-final-save-at-the-end-of-training-7"></a>Final save at the end of training</h1>
<p>final_checkpoint_path = os.path.join(checkpoint_dir, ‘model_final.pth’)<br>
torch.save({<br>
‘step’: max_steps,<br>
‘model_state_dict’: model.state_dict(),<br>
‘optimizer_state_dict’: optimizer.state_dict(),<br>
‘scheduler_state_dict’: scheduler.state_dict(),<br>
‘loss’: loss.item(),<br>
}, final_checkpoint_path)<br>
print(f"Final model saved to {final_checkpoint_path}")``</p>
<p><strong>Samplng code</strong><br>
``import torch<br>
from Prepro import Data_orig<br>
from Rnn_generator import RNNModel</p>
<p>def sample_and_infer(data, checkpoint_path, device, start_str, num_samples, block_size):<br>
“”"<br>
Samples data and performs inference using a pre-trained model.</p>
<pre><code>Parameters:
data (Data_orig): Preprocessed dataset.
checkpoint_path (str): Path to the model checkpoint.
device (torch.device): Device to use for the computation.
start_str (str): Starting string for the text generation.
num_samples (int): Number of characters to generate.
block_size (int): Size of the input sequence.

Returns:
str: Generated text.
"""
# Load pre-trained model
model = RNNModel(vocab_size=data.vocab_size, n_embd=10, n_hidden1=500, n_hidden2=200, block_size=block_size).to(device)
checkpoint = torch.load(checkpoint_path, map_location=device,weights_only=True)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)
model.eval()  # Set the model to evaluation mode

# Convert starting string to tensor
input_indices = torch.tensor([data.stoi[ch] for ch in start_str], device=device).unsqueeze(0)

# Generate text
generated = start_str
for _ in range(num_samples):
    with torch.no_grad():
        emb = model.embedding(input_indices)
        rnn_out, _ = model.rnn(emb)
        rnn_out = rnn_out.reshape(rnn_out.size(0), -1)
        logits = model.fc2(torch.tanh(model.fc1(rnn_out)))
        probs = torch.softmax(logits, dim=-1)
        next_index = torch.multinomial(probs, num_samples=1).item()

    generated += data.itos[next_index]
    input_indices = torch.cat((input_indices, torch.tensor([[next_index]], device=device)), dim=1)
    input_indices = input_indices[:, -block_size:]

return generated
</code></pre>
<h1><a name="p-459023-usage-example-8" class="anchor" href="#p-459023-usage-example-8"></a>Usage example</h1>
<p>dataset = Data_orig()<br>
checkpoint_path = r"C:\Users\Rick\Name_generator\checkpoints\model_step_90000.pth"<br>
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)</p>
<p>start_str = “a”<br>
num_samples = 10<br>
block_size = 3</p>
<p>generated_text = sample_and_infer(dataset, checkpoint_path, device, start_str, num_samples, block_size)<br>
print(“Generated Text:”)<br>
print(generated_text)``</p>
<p>Now when i execute the sampling code instead of loading the pretrained weights it again starts the training from beginning and saves the checkpoints.</p>

ptrblck: <aside class="quote no-group" data-username="Rajatavaa" data-post="1" data-topic="212724">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rajatavaa/48/73875_2.png" class="avatar"> Rajatavaa:</div>
<blockquote>
<p>Now when i execute the sampling code instead of loading the pretrained weights it again starts the training from beginning and saves the checkpoints.</p>
</blockquote>
</aside>
<p>Your code is unfortunately not properly formatted, but I assume you are using separate files for the training and sampling.<br>
It then seems that running <code>python sampling.py</code> is again training the model instead of executing the <code>sample_and_infer</code> function as you would expect?</p>
<p>If that’s the case, use <code>if __name__=="__main__"</code> guards in your scripts guarding the actual function calls, e.g.:</p>
<pre data-code-wrap="python"><code class="lang-python"># train.py

def MyModel(...):
    ...

def train(...):
    ...

if __name__=="__main__":
    train()
</code></pre>
<pre data-code-wrap="python"><code class="lang-python"># sampling.py

from Rnn_generator import RNNModel # because of the guard this will not execute any functions called in the __main__ block
...

def sample_and_infer(...):
    ...

if __name__=="__main__":
    sample_and_infer()
</code></pre>
<p>Without these guards the <code>import</code> statement will execute all global function calls and could thus start your training again.</p>

Rajatavaa: <p>This saved me…Thanks!<br>
I am still curious that why after specifying model.eval() the sampling script was starting the training process?</p>

ptrblck: <p>The Python import system will run the entire script which is defined outside of the <code>if __name__=="__main__"</code> method. I.e. if you are creating any objects, or are calling into functions, these will be executed during the import, which is why the import guard is important.<br>
Again, your code is not properly formatted, but as far as I can see you are initializing and training the model in the “global” space (i.e. it’s not wrapped in any methods).</p>
<p>Here is a small example:</p>
<pre data-code-wrap="python"><code class="lang-python"># fun.py
def fun():
    print("calling fun")
    
fun()
</code></pre>
<p>You can see that <code>fun.py</code> contains the <code>fun</code> function, which I want to import. However, the <code>fun()</code> function is also called in the global space and will thus be executed:</p>
<pre data-code-wrap="python"><code class="lang-python">import fun
calling fun
</code></pre>
<p>If you change it to:</p>
<pre data-code-wrap="python"><code class="lang-python"># fun.py
def fun():
    print("calling fun")

if __name__=="__main__":    
    fun()
</code></pre>
<p>the <code>import fun</code> will not print anything since the actual <code>fun()</code> call is now guarded and will only be executed if you execute this script directly:</p>
<pre data-code-wrap="python"><code class="lang-python">python fun.py 
calling fun
</code></pre>