CIM_Training/simcim_model_8classes_BNN_ginger.py

import torch
import numpy as np
import matplotlib
#matplotlib.use('TkAgg')
from matplotlib import pyplot as plt  
import cv2, os, sys
from torch.utils.data import Dataset

from torch.utils.data import ConcatDataset, DataLoader, Subset
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import DatasetFolder
from PIL import Image
from SimBinaryNetpytorch.models.binarized_modules import  CimSimConv2d
from SimBinaryNetpytorch.models.binarized_modules import  BinarizeConv2d, IdealCimConv2d
from BinaryNetpytorch.models.binarized_modules import  BinarizeConv2d as BConv2d
from BinaryNetpytorch.models.binarized_modules import  Binarize,HingeLoss
import seaborn as sns
import random
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
batch_size = 8
num_epoch = 150
LEARN_RATE = 0.001

seed = 333
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True


H = [32, 64, 128]
H = [16, 32, 64]

RANDOM_WEIGHT_PER_EPOCH = 10
RANDOM_RATE = 1


class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        conv = CimSimConv2d
        conv2 = BConv2d
        conv3 = IdealCimConv2d
        conv3 = BinarizeConv2d
        #conv = nn.Conv2d

        self.cnn_layers1 = nn.Sequential(
            
            #conv(in_channels=1, out_channels=128, kernel_size=7),
            #nn.BatchNorm2d(128),
            #nn.LeakyReLU(0.5),
            #conv(in_channels=128, out_channels=64, kernel_size=3),
            #nn.BatchNorm2d(64),
            #nn.LeakyReLU(0.5),
            
            #input_size(1,30,40)
            conv(1, H[0], 3, 1), #output_size(16,66,66)
            #nn.BatchNorm2d(128),
            nn.LeakyReLU(0.5),
            #nn.Dropout(0.2),
            nn.MaxPool2d(kernel_size = 2), #output_size(16,33,33)
        )
        self.cnn_layers2 = nn.Sequential(

            conv(H[0], H[1], 3, 1), #output_size(24,31,31)
            #nn.BatchNorm2d(64),
            nn.LeakyReLU(0.5),
            #nn.Dropout(0.2),
            nn.MaxPool2d(kernel_size = 2), #output_size(24,15,15)
        )

        self.cnn_layers3 = nn.Sequential(

            conv(H[1], H[2], 3, 1), #output_size(32,13,13)
            #nn.BatchNorm2d(32),
            nn.LeakyReLU(0.5),
            #nn.Dropout(0.2),
            nn.MaxPool2d(kernel_size = 2), #ouput_size(32,6,6)
            #nn.LogSoftmax(),
            #BinarizeConv2d(H[2], 8, (3,2), 1) #ouput_size(4,2,3) without max :(32,24,34)
            conv2(H[2], 8, (3,2), 1) #ouput_size(4,2,3) without max :(32,24,34)
            #conv(32, 8, (2,1), 1) #ouput_size(4,2,3) without max :(32,24,34)
                        
        )
        
            

    def forward(self, x):
        #print(x)
        #print("input",float(torch.min(x)),float(torch.max(x)))
        x = self.cnn_layers1(x)
        #print("layer1",float(torch.min(x)),float(torch.max(x)))
        #print(x)
        x = self.cnn_layers2(x)
        #print("layer2",float(torch.min(x)),float(torch.max(x)))
        x = self.cnn_layers3(x)
        #print("layer3",float(torch.min(x)),float(torch.max(x)))
        #print(x)
        #x = x.flatten(1)
        #x = self.fc_layers(x)
        #print(x.shape)
        x = x.view(x.size(0), -1)
        #print(x.shape)
        #x = nn.LogSoftmax(x)
        #print(x)

        return x 

class Cls_Dataset(Dataset):
    def __init__(self, input, target):
        self.input = input
        self.target = target
    def __getitem__(self, index):
        return self.input[index], self.target[index]
    def __len__(self):
        return len(self.input)

def Load_data(path, cls):
    data = []
    label = []
    f = 0
    for c in range(cls):
        t = 0
        img = []
        data_path = path + "/0"+str(c)+"/"
        for filename in os.listdir(data_path):
            # with open(data_path + "/" + filename, "rb") as f_in:
            
            tmp_input = cv2.imread(data_path + filename,cv2.IMREAD_UNCHANGED)
            #print(tmp_input)
            tmp_input = cv2.resize(tmp_input, (30,40), interpolation = cv2.INTER_AREA)
            tmp_input = tmp_input.astype(int)
            tmp_input = tmp_input//2
            tmp_input = tmp_input - 63
            tmp_input = np.where(tmp_input > 63, 63, tmp_input)
            tmp_input = tmp_input[:, :, np.newaxis]
            tmp_input = tmp_input.transpose(2,0,1)
        
            
                
            
            if img is not None:
                img.append(tmp_input)
            else:
                img = tmp_input    
        label_tmp = np.full((len(img),1),c)
        
        if f != 0:
            data = np.append(data,img,axis=0)
        else:
            data = img
        if f != 0:
            label = np.append(label,label_tmp,axis=0)
        else:
            label = label_tmp
        f = 1



    label = np.squeeze(label)

    label = np.array(label)
    #print(data)
    data = np.array(data)
    #print(label.shape)
    data = data.astype('float32')

    return data, label

def main():
    # test_set = DatasetFolder("./dataset/8cls_srcnnimg/test", loader=lambda x: Image.open(x), extensions="jpg", transform=test_tfm)
    # val_set =  DatasetFolder("./dataset/8cls_srcnnimg/train", loader=lambda x: Image.open(x), extensions="jpg", transform=test_tfm)
    global num_epoch

    train_data, train_label = Load_data("./dataset/8cls_grideye/train",8)
    train = Cls_Dataset(train_data, train_label)
    train_loader = DataLoader(
            train, batch_size=100, shuffle=True,
            num_workers=4, pin_memory=True, drop_last=True
        )

    test_data, test_label = Load_data("./dataset/8cls_grideye/test",8)

    test = Cls_Dataset(test_data, test_label)

    test_loader = DataLoader(
            test, batch_size=100, shuffle=True,
            num_workers=4, pin_memory=True, drop_last=True
        )

    val_data, val_label = Load_data("./dataset/8cls_grideye/val",8)

    val = Cls_Dataset(val_data, val_label)

    val_loader = DataLoader(
            val, batch_size=100, shuffle=True,
            num_workers=4, pin_memory=True, drop_last=True
        )




    save_path = 'models.ckpt'
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = Classifier().to(device)
    #optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    try:
        checkpoint = torch.load('training_state_pre.bin')
        model.load_state_dict(checkpoint['state_dict'])
        #optimizer.load_state_dict(checkpoint['optimizer'])
        #model.load_state_dict(torch.load('models_pre.ckpt'))
        num_epoch = 150
    except:
        print('cannot read from pretrained model.')
    criterion = nn.CrossEntropyLoss()
    best_accuracy = 0.0
    last_save = 0
    for epoch in range(num_epoch):
        model.train()
        running_loss = 0.0
        total = 0
        correct = 0
        for  i, data in enumerate(train_loader):
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()

            outputs = model(inputs)
            #print(outputs.shape)

            loss = criterion(outputs, labels)
            loss.backward()
            #print(model)
            #print(model.cnn_layers3[3].weight.grad)


            for p in list(model.parameters()):
                if hasattr(p,'org'):
                    p.data.copy_(p.org) 

            optimizer.step()

            for p in list(model.parameters()):
                if hasattr(p,'org'):
                    p.org.copy_(p.data.clamp_(-1,1))
            
            running_loss += loss.item()
            total += labels.size(0)
            _,predicted = torch.max(outputs.data,1)
            #print(predicted)
            #print("label",labels)
            correct += (predicted == labels).sum().item()
        train_acc = correct / total

        print(f"[ Train | {epoch + 1:03d}/{num_epoch:03d} ] loss = {running_loss:.5f}, acc = {train_acc:.5f}")

        model.eval()
        with torch.no_grad():
            correct = 0
            total = 0
            for  i, data in enumerate(val_loader):
                inputs, labels = data
                inputs = inputs.to(device)
                labels = labels.to(device)
                outputs = model(inputs)
                _,predicted = torch.max(outputs.data,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
            val_acc = correct / total
        
        if val_acc > best_accuracy:
            last_save = epoch
            best_accuracy = val_acc
            torch.save(model.state_dict(), save_path)
            state = {'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
            torch.save(state, 'training_state.bin')
            print("Save Model")
        
        print(f"[ Val | {epoch + 1:03d}/{num_epoch:03d} ]  acc = {val_acc:.5f}")
        

    model = Classifier().to(device)
    model.load_state_dict(torch.load(save_path))
    model.eval()
    stat = np.zeros((8,8))
    with torch.no_grad():
        correct = 0
        total = 0
        print(model)
        for i, data in enumerate(test_loader):
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            #print(outputs.data)
            _,predicted = torch.max(outputs.data,1)
            #print(predicted)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            #print(labels)
            #print(outputs.size())
            

        print('Test Accuracy:{} %'.format((correct / total) * 100))
        print('Save at', last_save+1)
if __name__ == '__main__':
    main()