I get something wrong when use model.train() and model.eval() on pytorch

Viewed 1084

I have prepare features and their labels as blow; I want to build a model which is constructed by transformers' encoder and then add a linear layer to predict a value. but I got some error when I use the model to predict after its training. At first I run below code:

import torch
from torch import nn

features = torch.rand(bach_size, channels, lenght)
labels = torch.rand(batch_size)

class TransformerModel(nn.Module):

  def __init__(self):
    super(TransformerModel, self).__init__()
    encoder_layer = nn.TransformerEncoderLayer(d_model=8, nhead=8, dropout=0.5)
    self.transformer_encoder = nn.TransformerEncoder(encoder_layer, 6)
    self.decoder = nn.Linear(40, 1)

  def forward(self, src):
    encoded = self.transformer_encoder(src.transpose(1, 0)).transpose(1, 0)
    pred = self.decoder(encoded.reshape(encoded.shape[0], -1))
    return pred


model = TransformerModel()
criterion = nn.MSELoss()
lr = 0.3 # learning rate
optimizer = torch.optim.Adam(model.parameters(), lr=lr)


def train():

  model.train() # Turn on the train mode
  optimizer.zero_grad()

  output = model(features)

  loss = criterion(output.view(-1, 1), labels.view(-1, 1))
  loss.backward()
  torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
  optimizer.step()

  return loss.item()



for _ in range(100):
  train()

After that, I predict features by the below codes:

model.eval()
output = model(features)

I get all values of 'output' are the same, and if use 'model.train()', the 'output' seems Ok; so what is the problem? or the model was built wrong?

0 Answers
Related