Wanting to make a simple age estimator

username: Jia_Ming
cooked: <p>I am planning to make an age detector with 10 classes, each of them has range from 2-6 years old, 7-12 years old and so on.<br>
I use pretrained model from resnet18. During training, I did not freeze the layers, instead, I just let it update all the parameters. The loss function I am using is cross-entropy, with Adam optimizer and lr-scheduler. The datasets contain 10000 images for training and about 3000 images for validation.</p>
<p>The problem that keeping me stuck is that although the training loss seems able to continuous decreased, but the validation loss is not. No matter how I change the hyper-parameters, It will drop until minimum 33 and bounce back to 60+, which is the initial value when I started training. The accuracy for validation is only at most 46%</p>
<p>This is my code. Please have a look at what is causing this problem. Or is there any problem with datasets? Such as not enough datasets for training etc?</p>
<pre><code class="lang-auto">    if train_mode:
        train_loader, test_loader = data_load()
        for iteration in range(iteration_start, epoch):
            time_start = time.time()
            for current_mode in ['train', 'valid']:
                total_loss = 0
                total_accuracy = 0
                if current_mode == 'train':
                    loader = train_loader
                    model.train()
                else:
                    loader = test_loader
                    model.eval()

                for batch in loader:
                    images, labels = batch
                    images = images.to(device)
                    labels = labels.to(device)

                    with torch.set_grad_enabled(current_mode == 'train'):
                        output = model(images)
                        loss = loss_function(output, labels)
                        total_loss += loss.item()
                        if current_mode == 'train':
                            optimizer.zero_grad()
                            loss.backward()
                            optimizer.step()
                        else:
                            accuracy = calculate_accuracy(output, labels)
                            total_accuracy += accuracy

                record_loss(total_loss, current_mode, iteration)
                if current_mode == 'valid':
                    scheduler.step(total_loss)  # total_loss or total_accuracy, based on which you want to enhance
                    record_accuracy(total_accuracy / len(loader), iteration)
</code></pre>
<p>This is my model</p>
<pre><code class="lang-auto">def custom_model():
    my_model = models.resnet18(pretrained=True)
    # for paras in my_model.parameters():
    #     paras.requires_grad = False

    num_fc_layer = my_model.fc.in_features
    custom_fc_layers = nn.Sequential(
        nn.BatchNorm1d(num_fc_layer),
        nn.Dropout(0.5),
        nn.Linear(num_fc_layer, num_class)
    )
    my_model.fc = custom_fc_layers

    return my_model

</code></pre>
<p>and this is my the state of optimizer and scheduler</p>
<pre><code class="lang-auto">if __name__ == '__main__':
    model = custom_model()
    model.to(device)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=1, verbose=True)
    loss_function = nn.CrossEntropyLoss()
    main()
</code></pre>