continuation of previous post.
pz_train_resnet.py : the script for training
#This is a standalone script to train resnet for debugging
import argparse
import os
import random
import types

import torch
import torch.nn.functional as F

import numpy as np

from pz_utils import *

## ###################### TRAIN ##################
def train(args, model, device, train_loader, transforms, optimizer, epoch,
          weights_class=None):

    """
    Training phase 
    Minimize Cross Entropy Loss
    """

    assert transforms

    # switch network layers to Training mode
    model.train()

    n_bins = model.n_bins # the last number of neurons

    # JEC Todo : fix them both for training and testing
    z_min = 0.0
    z_max = 1.0


    train_loss = 0
    # scans the batches
    Nbatch = len(train_loader)

    for i_batch, sample_batched in enumerate(train_loader):
        img_batch = sample_batched['image']   #input image
        ebv_batch = sample_batched['ebv']     #input redenning
        z_batch   = sample_batched['z']       #target

        batch_size = len(img_batch)

        transf_size = len(transforms)
        new_img_batch = torch.zeros_like(img_batch).permute(0, 3, 1, 2)

        #for CrossEntropyLoss no hot-vector
        new_z_batch = torch.zeros(batch_size,dtype=torch.long)

        for i in range(batch_size):
            # transform the images
            img = img_batch[i].numpy()

            for it in range(transf_size):
                img = transforms[it](img)

            new_img_batch[i] = img

            # transform the redshift in bin number
            z = (z_batch[i] - z_min) / (z_max - z_min)    # z \in 0..1                => z est reel
            z = max(0, min(n_bins - 1, int(z * n_bins)))  # z \in {0,1,.., n_bins-1}  => z est entier
            new_z_batch[i] = z

        # send the inputs and target to the device
        new_img_batch, ebv_batch, new_z_batch = new_img_batch.to(device), \
                                            ebv_batch.to(device), \
                                            new_z_batch.to(device)


        # reset the gradiants
        optimizer.zero_grad()
        # Feedforward
        output = model(new_img_batch, ebv_batch)


        # the loss
        loss = F.cross_entropy(output,new_z_batch, weight=weights_class)
        train_loss += loss.item()
        # backprop to compute the gradients
        loss.backward()
        # perform an optimizer step to modify the weights
        optimizer.step()

        # some debug
        if i_batch % (Nbatch//10) == 0:
             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                 epoch, i_batch * batch_size, Nbatch*batch_size,
                 100. * i_batch / Nbatch, loss.item()))


    # return loss.item()
    return train_loss/Nbatch

## ################### TEST ###################

def test(args, model, device, test_loader, transforms, epoch,
         weights_class=None):
    """
    Testing phase.
    
    """
    # switch network layers to Testing mode
    model.eval()
    n_bins = model.n_bins # the last number of neurons
    # JEC Todo : fix them both for training and testing
    z_min = 0.0
    z_max = 1.0
    largeur_bin=(z_max-z_min)/n_bins

    test_loss = 0
    correct   = 0
    #    # turn off the computation of the gradient for all tensors
    Nbatch = len(test_loader)
    with torch.no_grad():
        # scans the batches
        for i_batch, sample_batched in enumerate(test_loader):
            img_batch = sample_batched['image']  # input image
            ebv_batch = sample_batched['ebv']    # input redenning
            z_batch   = sample_batched['z']      # target

            # transform the images
            batch_size = len(img_batch)
            transf_size = len(transforms)
            new_img_batch = torch.zeros_like(img_batch).permute(0, 3, 1, 2)
            # for CrossEntropyLoss no hot-vector needed
            new_z_batch = torch.zeros(batch_size, dtype=torch.long)

            for i in range(batch_size):
                # transform the images
                img = img_batch[i].numpy()

                for it in range(transf_size):
                    img = transforms[it](img)

                new_img_batch[i] = img

                # transform by hand to hot vector the redshift
                z = (z_batch[i] - z_min) / (z_max - z_min)    # z \in 0..1                => z est reel
                z = max(0, min(n_bins - 1, int(z * n_bins)))  # z \in {0,1,.., n_bins-1}  => z est entier
                new_z_batch[i] = z

            # send the inputs and target to the device
            new_img_batch, ebv_batch, new_z_batch = new_img_batch.to(device), \
                                                ebv_batch.to(device), \
                                                new_z_batch.to(device)


            # Feedforward
            output = model(new_img_batch, ebv_batch)


            # the loss (w/o the mean average)
            test_loss += F.cross_entropy(output, new_z_batch,
                                         weight=weights_class).item()
            #
            proba = F.softmax(output, dim=1)
            pred  = proba.argmax(dim=1)
            correct_current = sum([1 if y1==y2 else 0 for y1, y2 in zip(pred, new_z_batch)])
            correct += correct_current


    test_loss /= Nbatch
    accuracy = 100. * correct / (Nbatch*batch_size)

    return {'loss': test_loss, 'acc': accuracy}

## ################### Reference TEST to be saved ###################
def test_ref(args, model, device, sample_batched, transforms, epoch):
    """
    Testing  reference

    """
    # switch network layers to Testing mode
    model.eval()
    n_bins = model.n_bins # the last number of neurons
    # JEC Todo : fix them both for training and testing
    z_min = 0.0
    z_max = 1.0
    largeur_bin=(z_max-z_min)/n_bins


    with torch.no_grad():

        img_batch = sample_batched['image']  # input image
        ebv_batch = sample_batched['ebv']    # input redenning
        z_batch   = sample_batched['z']      # target

        # transform the images
        batch_size = len(img_batch)
        transf_size = len(transforms)
        new_img_batch = torch.zeros_like(img_batch).permute(0, 3, 1, 2)
        # for CrossEntropyLoss no hot-vector needed
        new_z_batch = torch.zeros(batch_size, dtype=torch.long)

        for i in range(batch_size):
            # transform the images
            img = img_batch[i].numpy()

            for it in range(transf_size):
                img = transforms[it](img)

            new_img_batch[i] = img

            # transform by hand to hot vector the redshift
            z = (z_batch[i] - z_min) / (z_max - z_min)    # z \in 0..1                => z est reel
            z = max(0, min(n_bins - 1, int(z * n_bins)))  # z \in {0,1,.., n_bins-1}  => z est entier
            new_z_batch[i] = z

        # send the inputs and target to the device
        new_img_batch, ebv_batch, new_z_batch = new_img_batch.to(device), \
                                            ebv_batch.to(device), \
                                            new_z_batch.to(device)


        # Feedforward
        output = model(new_img_batch, ebv_batch)


    # return the output reference
    return {'ref_in': new_img_batch, 'ref_ebv': ebv_batch, 'ref_out': output}


## ############################### MAIN ########################

def main():
    args = types.SimpleNamespace(
        fake_data = False,         # True means random inputs, False means either REAL or SYNTHISED data
        use_weighted_class = False,  # to reweight the class
        weights_file = './synth_weight.npy',
        network = "resnet20",      # alternative : "resnet20" or "convnet"
        deterministic = False,     #    torch.backends.cudnn.deterministic
        benchmark     = True,      #    torch.backends.cudnn.benchmark
        data_augmentation = True, # use or not data augmentation
        resume = False,            # resume a session
        checkpoint_file = "dbg-resnet20_cudnn_False_True_data_True_19.pth",
        history_loss_cpt_file = "dbg-history-resnet20_cudnn_False_True_data_True_19.npy",
        Nepochs         = 20, 
        batch_size      = 128,
        test_batch_size = 128,
        lr_init = 0.01,
        weight_decay = 0.0,
        momentum = 0.9,
        lr_decay = 0.1,
        no_cuda = False,
        root_file = "./",
        # data used if fake_data=False
# Real data
#        train_file     = './data/train100k.npz',
#        test_file      = './data/test100k.npz',
#        test_ref_file  = './data/test10k.npz'
# Synthetised data
        train_file     = './data/train_synth_128k.npz',
        test_file      = './data/test_synth_128k.npz',
        test_ref_file  = './data/test_synth_10k.npz'
    )
    
    args.checkpoint_file = args.root_file + args.checkpoint_file
    args.history_loss_cpt_file = args.root_file + args.history_loss_cpt_file


    print("\n### Training model ###")
    print("> Parameters:")
    for p, v in zip(args.__dict__.keys(), args.__dict__.values()):
        print('\t{}: {}'.format(p, v))

    use_cuda = not args.no_cuda and torch.cuda.is_available()

    device = torch.device("cuda" if use_cuda else "cpu")
    print("Use device....: ",device)

    # Pin_memory speed up the CPU->GPU transfert for NVidia GPU
    kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}


    weights_class = None # if reweighting is needed 
    
    if args.fake_data: # Use fake data train/test/test_ref
        ######################################        
        # Train data
        Ntrain = 10**3
        train_loader = []
        for i in range(Ntrain):
            img_train = torch.zeros(args.batch_size,64,64,5).uniform_(-1.,10.) # batch_sizexHxWxC
            ebv_train = torch.zeros(args.batch_size,1).exponential_(0.05)
            z_train   = torch.zeros(args.batch_size,1).normal_(0.1,0.03).clamp_(0,0.3)
            train_loader.append({'image':img_train,'ebv':ebv_train,'z':z_train})


        # Test data
        Ntest = 10**3
        test_loader = []
        for i in range(Ntest):
            img_test = torch.zeros(args.batch_size,64,64,5).uniform_(-1.,10.) # batch_sizexHxWxC
            ebv_test = torch.zeros(args.batch_size,1).exponential_(0.05)
            z_test   = torch.zeros(args.batch_size,1).normal_(0.1,0.03).clamp_(0,0.3)
            test_loader.append({'image':img_test,'ebv':ebv_test,'z':z_test})


        # Reference test
        test_ref_loader = []
        img_test_ref = torch.ones(10,64,64,5) # NxHxWxC
        ebv_test_ref = torch.ones(10,1)
        z_test_ref   = torch.ones(10,1)
        test_ref_loader.append({'image':img_test_ref,'ebv':ebv_test_ref,'z':z_test_ref})

    else:  # Use real/synthetized data train/test/test_ref
        ######################################
        if args.use_weighted_class :
            weights_class = np.load(args.weights_file)
            weights_class = torch.from_numpy(weights_class).to(device=device,dtype=torch.float)
        
        # Train data
        train_dataset = DatasetPz(args.train_file, transform=None)
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.batch_size,drop_last=True,
            shuffle=True, **kwargs)
        print("train_loader length=",len(train_loader))#, " WARNING: NO SHUFFLE")

        # Test data
        test_dataset = DatasetPz(args.test_file, transform=None)
        test_loader = torch.utils.data.DataLoader(
            test_dataset,
            batch_size=args.test_batch_size, drop_last=True,
            shuffle=False, **kwargs)
        print("test_loader length=",len(test_loader))
        # Reference data
        test_ref_dataset = DatasetPz(args.test_ref_file, transform=None)
        test_ref_loader = torch.utils.data.DataLoader(
            test_ref_dataset,
            batch_size=10, drop_last=True,
            shuffle=False, **kwargs)
        print("test_ref_loader length=",len(test_ref_loader))
        testref_loader_iterator = iter(test_ref_loader)
        sample_test_ref = next(testref_loader_iterator) # once for all



    # The Network
    if args.network == "resnet20":
        model = resnet20()
    else:
        args.network = "convnet"
        model = NetCNNRed(5)
    model.to(device)


    #Optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr_init,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay,
                                nesterov=False)

    #Scheduler
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
                                                     factor=args.lr_decay, patience=5,
                                                     verbose=True)

    





    #Optional data augmentation for test set
    if args.data_augmentation:
        train_transforms = [RandomApplyPz([[flipH,flipV,identity],
                                           [rot90,rot180,rot270,identity]]),
                            ToTensorPz()]
    else:
        train_transforms = [ToTensorPz()]

    # No data augmentation for test set
    test_transforms = [ToTensorPz()]

    # Loop on the epochs
    train_loss_history = []
    test_loss_history  = []

    # load a previous session if required
    start_epoch = 0
    if args.resume :
        # load checkpoint of model/scheduler/optimizer
        if os.path.isfile(args.checkpoint_file):
            print("=> loading checkpoint '{}'".format(args.checkpoint_file))
            checkpoint = torch.load(args.checkpoint_file)
            # the first epoch for the new training
            start_epoch = checkpoint['epoch']
            # model update state
            model.load_state_dict(checkpoint['model_state_dict'])
            # scheduler update state
            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
            # optizimer update state
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        else:
            print("=> FATAL no  checkpoint '{}'".format(args.checkpoint_file))
            return

        #load previous history of losses
        if os.path.isfile(args.history_loss_cpt_file):
            loss_history = np.load(args.history_loss_cpt_file)
            train_loss_history = loss_history[0].tolist()
            test_loss_history  = loss_history[1].tolist()
        else:
            print("=> FATAL no history loss checkpoint '{}'".format(args.history_loss_cpt_file))
            return
    else:
        print("=> no checkpoints then Go as fresh start")


    #seeds manual init
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed_all(0)
    random.seed(0)

    # Cuda Cudnn
    torch.backends.cudnn.deterministic = args.deterministic
    torch.backends.cudnn.benchmark = args.benchmark

    for epoch in range(start_epoch, start_epoch+args.Nepochs):

        print("process epoch[",epoch,"]: LR = ",end='')
        for param_group in optimizer.param_groups:
            print(param_group['lr'])

        # training
        train_loss = train(args, model, device, train_loader,
                           train_transforms, optimizer, epoch, weights_class)

        # test
        test_loss = test(args, model, device, test_loader,
                         test_transforms, epoch, weights_class)
        

        # bookkeeping
        train_loss_history.append(train_loss)
        test_loss_history.append(test_loss['loss'])

        print(f"Epoch {epoch}: Train loss {train_loss:.6f}, Test loss {test_loss['loss']:.6f}")

        #step of the scheduler
        scheduler.step(test_loss['loss'])

        # Ref test
        test_ref_out = test_ref(args, model, device,  sample_test_ref,
                                test_transforms, epoch)

        

        #save model state and reference test in/out
        state = {
            'epoch': epoch + 1,
            'model_state_dict': model.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            **test_ref_out
            }

        torch.save(state,args.root_file+"/dbg-"+args.network+"_cudnn_"+str(args.deterministic)+"_"+str(args.benchmark)+"_data_"+str(args.data_augmentation)+"_"+str(epoch)+".pth")
        # save intermediate history
        np.save(args.root_file+"/dbg-history-"+args.network+"_cudnn_"+str(args.deterministic)+"_"+str(args.benchmark)+"_data_"+str(args.data_augmentation)+"_"+str(epoch)+".npy",
                np.array((train_loss_history,test_loss_history))
                )

    #loop on epochs done
    print('End of job!')


################################
if __name__ == '__main__':
  main()

pz_debug_bn.py: a script to look at loss after a relaoding of a checkpoint (infliuence of Cuda backend)
#This is a standalone script to train resnet for debugging

import argparse
import os
import random
import types

import numpy as np

from pz_utils import *

############## TRAIN BATCH ################
def process(model,sample_batched,transforms,train=None,debug=False):
    #We load the model checkpoint so no more training
    model.eval()

    n_bins = model.n_bins # the last number of neurons

    z_min = 0.0
    z_max = 1.0

    #    # turn off the computation of the gradient for all tensors
    with torch.no_grad():

        #get next data
        img_batch = sample_batched['image']   #input image
        ebv_batch = sample_batched['ebv']     #input redenning
        z_batch   = sample_batched['z']       #target

        batch_size  = len(img_batch)
        transf_size = len(transforms)

        new_img_batch = torch.zeros_like(img_batch).permute(0, 3, 1, 2)

        #for CrossEntropyLoss no hot-vector
        new_z_batch = torch.zeros(batch_size,dtype=torch.long)

        for i in range(batch_size):
            # transform the images
            img = img_batch[i].numpy()

            for it in range(transf_size):
                img = transforms[it](img)

            new_img_batch[i] = img

            # transform the redshift in bin number
            z = (z_batch[i] - z_min) / (z_max - z_min)    # z \in 0..1                => z est reel
            z = max(0, min(n_bins - 1, int(z * n_bins)))  # z \in {0,1,.., n_bins-1}  => z est entier
            new_z_batch[i] = z

        # send the inputs and target to the device
        new_img_batch, ebv_batch, new_z_batch = new_img_batch.to(device), \
                                                ebv_batch.to(device), \
                                                new_z_batch.to(device)


        # Feedforward
        output = model(new_img_batch, ebv_batch)
        print("ebv_batch: ",ebv_batch[:5])
        print("new_img_batch:\n",new_img_batch[:5])
        print("output:\n",output[:5])

        # the loss
        loss = F.cross_entropy(output,new_z_batch)

    return loss.item()



############## MAIN ################

#
# set manual seeds 
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
random.seed(0)

#Cuda 
#torch.backends.cudnn.enabled = False
#torch.backends.cudnn.deterministic = False
#torch.backends.cudnn.benchmark = True


use_cuda = torch.cuda.is_available()

device = torch.device("cuda" if use_cuda else "cpu")
print("Use device....: ",device)

# Pin_memory speed up the CPU->GPU transfert for NVidia GPU
kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}


#model
model = resnet20()
#model = NetCNNRed(5)
model.to(device)

#load model parameters
root_file = "./"
checkpoint_file = "dbg-resnet20_cudnn_False_True_data_True_19.pth.SAV"


checkpoint_file = root_file + checkpoint_file
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['model_state_dict'])


#x-check model inference form reference point
ref_input_img = checkpoint['ref_in']
ref_input_ebv = checkpoint['ref_ebv']
ref_output    = checkpoint['ref_out']
model.eval()
with torch.no_grad():
    output = model(ref_input_img,ref_input_ebv)
# Compare
diff = (output - ref_output).abs().max().item()
print('max abs error: ', diff)
if diff>0:
    print('ref_out\n',ref_output,'\n',F.softmax(ref_output,dim=1))
    print('new_out\n',output,'\n',F.softmax(output,dim=1))


With all that,

you can generate synthesized data
train resnet20 (and a simpler convnet) in different conditions
a) fixing Cuda backend
b) using or not pure random data
c) use unweighted classes or reweight the classes to counter-balance the class label distribution given by the ‘z’ varaible distribution binned in a hot-vector
during the training every epoch I save a ‘pth’-file for the model/optimizer/scheduler… checkpoint and a ‘npy’-file to track the train/test loss. You can plot them with the following script

import numpy as np
import matplotlib.pyplot as plt
import argparse

def main():

    parser = argparse.ArgumentParser(description='loss-plot')
    parser.add_argument('--file', type=str, required=True)
    parser.add_argument('--tag', type=str, required=True)
    
    args = parser.parse_args()

    data = np.load(args.file,allow_pickle=True)
    
    train_loss = data[0]
    test_loss = data[1]
    plt.plot(train_loss, label="train")
    plt.plot(test_loss, label="test")
    plt.xlabel('epoch')
    plt.ylabel('Cross Entropy loss')
    plt.legend()
    plt.grid()
    plt.savefig('history_'+args.tag+'.png')
        
    plt.show()

if __name__ == '__main__':
  main()


To my numerical experiments: the test loss tends to be hieratic with the un-reweighted classes synthesized data but this is not the case for real data (ie. reweighting as no effect on test loss hieratic behaviour), certainly you will ask me questions before.