training.py

import torch
from orm_model import CTC_CNN, default_model_params
from orm_dataset import CTC_PriMuS
import os
import time
import cv2
from PIL import Image
import matplotlib as plt

def decode(target):
    decoded = []
    decoded_2 = []
    prev = 0
    for note in target:
        if note == prev:
            continue
        else:
            decoded.append(note)
            prev = note
    for note in decoded:
        if note != 0:
            decoded_2.append(note)
    return decoded_2

data_dir = './package_aa'
dict_path = './vocabulary_semantic.txt'
train_loss=val_loss=[]
img_height = 128
learning_rate = 0.04
num_epochs = 500
batch_size = 16
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

corpus_list = os.listdir(data_dir)[0:700]
primus = CTC_PriMuS(data_dir, corpus_list, dict_path, True, val_split=0.1)
params = default_model_params(img_height, primus.vocabulary_size)
#print(primus.nextBatch(["targets"])[0][0][50:100])

# model
model = CTC_CNN(img_height, primus.vocabulary_size).to(device)
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
#scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
# loss
loss_func = torch.nn.CTCLoss()
loss_log = []

#batch = primus.nextBatch(params)
#inputs = torch.tensor(batch['inputs']).to(device)
#targets = torch.tensor(batch['targets']).to(device)
#inputs=torch.transpose(inputs,1,2)
   
#inputs=torch.transpose(inputs,1,3)
            
# loop
start_time = time.time()
for epoch in range(num_epochs):
    print(f'Epoch {epoch}/{num_epochs - 1}')
    print('-' * 10)

    epoch_time = time.time()

    # Each epoch has a training and validation phase
    for phase in ['train', 'val']:
        if phase == 'train':
            model.train()  # Set model to training mode
        else:
            val_i = 0
            model.eval()

        running_loss = 0.0
        running_corrects = 0

        # Iterate over data.
        while True:
            if phase == 'train':
                 batch = primus.nextBatch(params)
                 #print(primus.current_idx)
                 #img = Image.fromarray(batch['inputs'][0][:][:][:]*255,'RGB')
           
                  #img.show()
            
                 inputs = torch.tensor(batch['inputs']).to(device)
                 targets = torch.tensor(batch['targets']).to(device)
                #input_image=cv2.imread("C:\\Users\\psiml\\Downloads\\MusicDetector\\package_aa\\000051650-1_1_1\\000051650-1_1_1.png")
                #inputs=torch.tensor([input_image]).to(device)
                #output_line="clef-G2	keySignature-EbM	timeSignature-3/4	note-Bb5_quarter	note-Eb5_eighth	note-Bb5_eighth	note-C6_eighth	note-Bb5_eighth	barline	note-Ab5_eighth	note-Ab5_eighth	rest-sixteenth	note-Ab5_sixteenth	note-G5_sixteenth	note-Ab5_sixteenth	note-Bb5_sixteenth	note-Ab5_sixteenth	note-G5_sixteenth	note-Ab5_sixteenth	barline	"
                #targets =  torch.tensor([output_line]).to(device)
            else:
                batch = primus.getValidation(params)
                inputs = torch.tensor(batch['inputs'][val_i:val_i+params['batch_size']]).to(device)
                targets = torch.tensor(batch['targets'][val_i:val_i+params['batch_size']]).to(device)

            #inputs = inputs.view(inputs.shape[0], inputs.shape[3], inputs.shape[1], inputs.shape[2])
            inputs=torch.transpose(inputs,1,2)
            #print(inputs.shape)
            inputs=torch.transpose(inputs,1,3)
            #print(inputs.shape)


            # zero the parameter gradients
            optimizer.zero_grad()
                
            # forward
            # track history if only in train
            with torch.set_grad_enabled(phase == 'train'):
                outputs = model(inputs)
                outputs_transposed = outputs.permute(1,0,2)
                
                output_lengths = [len(x) for x in outputs]
                target_lenghts = [len(x) for x in targets]
                #print(output_lengths)
                #print(target_lenghts)
                #print(torch.argmax(outputs[5], dim=1))
                #print(targets[5])

                decoded_out = [torch.argmax(out, dim=1) for out in outputs]
                
                decoded_target = [target for target in targets]
                
                correct = 0
                total = 0

                for out, target in zip(decoded_out, decoded_target):
                    total += len(target)
                    for o, t in zip(out, target):
                        if o == t:
                            correct += 1
                
                print(decoded_out[0])
                print(decoded_target[0])
                #print(outputs.size)

                running_corrects = correct / total

                loss = loss_func(outputs_transposed, targets, output_lengths, target_lenghts)

                # backward + optimize only if in training phase
                if phase == 'train':
                    loss.backward()
                    #print(model.conv1.weight.grad)
                    optimizer.step()

                # statistics TODO: pitaj
                running_loss += loss.item() * inputs.size(0)
                #running_corrects += torch.sum(preds == labels.data).item()
            if phase == 'train':
                if primus.current_idx == 0:
                    break
            else:
                val_i += params['batch_size']
                if val_i >= len(primus.validation_list):
                    break

        if phase == 'train':
            epoch_loss = running_loss/ len(primus.training_list)
            #loss_log.append([epoch_loss])
            train_loss.append(epoch_loss)
        else:
            epoch_loss = running_loss / len(primus.validation_list)
            #loss_log[-1].append(epoch_loss)
            val_loss.append(epoch_loss)
        #epoch_acc = float(running_corrects) / dataset_sizes[phase]

        print(f'{phase} Loss: {epoch_loss:.4f}')#' Acc: {epoch_acc:.4f}')
        if phase == 'train':
            print(f'{phase} Acc: {running_corrects/16:.4f}')
        else:
            print(f'{phase} Acc: {running_corrects/16:.4f}')

            
        #metrics[phase+"_loss"].append(epoch_loss)
        #metrics[phase+"_acc"].append(epoch_acc)
        '''    
            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())'''

        print()
    #scheduler.step()
    torch.save(model, './model1')
    print(time.time() - epoch_time)

time_elapsed = time.time() - start_time
print(f'Training complete in {(time_elapsed // 60):.0f}m {time_elapsed % 60:.0f}s')
plt.plot(train_loss,np.linspace(0,len(train_loss),len(train_loss)),color='r')
plt.show()
plt.plot(val_loss,np.linspace(0,len(val_loss),len(val_loss)),color='g')
plt.show()
#print(loss_log)
#print('Bestval Acc: {best_acc:4f}')