import torch import torch.nn as nn from torch.nn import functional as F import torch.optim as optim from torch.optim import lr_scheduler import torch.backends.cudnn as cudnn import numpy as np import torchvision from torchvision import datasets, models, transforms import matplotlib.pyplot as plt import time import os from PIL import Image from tempfile import TemporaryDirectory import streamlit as st cudnn.benchmark = True plt.ion() # interactive mode class classifier(): def __init__(self): self.data_transforms = None self.data_dir = None self.image_datasets = None self.dataloaders = None self.dataset_sizes = None self.class_names = None self.device = None self.num_classes = None def data_loader(self,path,batch_size=4): # Data augmentation and normalization for training # Just normalization for validation self.data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'test': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } self.data_dir = path self.image_datasets = {x: datasets.ImageFolder(os.path.join(self.data_dir, x), self.data_transforms[x]) for x in ['train', 'val','test']} self.dataloaders = {x: torch.utils.data.DataLoader(self.image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'val','test']} self.dataset_sizes = {x: len(self.image_datasets[x]) for x in ['train', 'val','test']} self.class_names = self.image_datasets['train'].classes self.num_classes = len(self.class_names) self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def train(self,model, criterion, optimizer, scheduler, num_epochs=25): since = time.time() # Create a temporary directory to save training checkpoints with TemporaryDirectory() as tempdir: best_model_params_path = os.path.join(tempdir, 'best_model_params.pt') torch.save(model.state_dict(), best_model_params_path) best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch+1}/{num_epochs}') print('-' * 10) st.sidebar.subheader(f':blue[Epoch {epoch+1}/{num_epochs}]', divider='blue') # st.sidebar.code('-' * 10) # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': model.train() # Set model to training mode else: model.eval() # Set model to evaluate mode running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in self.dataloaders[phase]: inputs = inputs.to(self.device) labels = labels.to(self.device) # zero the parameter gradients optimizer.zero_grad() # forward # track history if only in train with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # backward + optimize only if in training phase if phase == 'train': loss.backward() optimizer.step() # statistics running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) if phase == 'train': scheduler.step() epoch_loss = running_loss / self.dataset_sizes[phase] epoch_acc = running_corrects.double() / self.dataset_sizes[phase] print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') st.sidebar.caption(f':blue[{phase[0].upper() + phase[1:]} Loss:] {epoch_loss:.4f} :blue[ Accuracy:] {epoch_acc:.4f}') # deep copy the model if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc torch.save(model.state_dict(), best_model_params_path) print() time_elapsed = time.time() - since print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s') print(f'Best val Accuracy: {best_acc:4f}') st.sidebar.caption(f':green[Training complete in] {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s') st.sidebar.subheader(f':blue[Best val Accuracy:] {best_acc:4f}') # load best model weights model.load_state_dict(torch.load(best_model_params_path)) return model def train_model(self,model_name,epochs): num_classes = self.num_classes if model_name == 'EfficientNet_B0': model = models.efficientnet_b0(pretrained=True) model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) # model.classifier[1].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'EfficientNet_B1': model = models.efficientnet_b1(pretrained=True) model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) # model.classifier[1].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'MnasNet0_5': model = models.mnasnet0_5(pretrained=True) model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) # model.classifier[1].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'MnasNet0_75': model = models.mnasnet0_75(pretrained=True) model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) # model.classifier[1].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'MnasNet1_0': model = models.mnasnet1_0(pretrained=True) model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) # model.classifier[1].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'MobileNet_v2': model = models.mobilenet_v2(pretrained=True) model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) # model.classifier[1].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'MobileNet_v3_small': model = models.mobilenet_v3_small(pretrained=True) model.classifier[3] = nn.Linear(model.classifier[3].in_features, num_classes) # model.classifier[3].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[3].parameters(), lr=0.001, momentum=0.9) elif model_name == 'MobileNet_v3_large': model = models.mobilenet_v3_large(pretrained=True) model.classifier[3] = nn.Linear(model.classifier[3].in_features, num_classes) # model.classifier[3].out_features = num_classes optimizer = torch.optim.SGD(model.classifier[3].parameters(), lr=0.001, momentum=0.9) elif model_name == 'RegNet_y_400mf': model = models.regnet_y_400mf(pretrained=True) model.fc = nn.Linear(model.fc.in_features, num_classes) # model.fc.out_features = num_classes optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9) elif model_name == 'ShuffleNet_v2_x0_5': model = models.shufflenet_v2_x0_5(pretrained=True) model.fc = nn.Linear(model.fc.in_features, num_classes) # model.fc.out_features = num_classes optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9) elif model_name == 'ShuffleNet_v2_x1_0': model = models.shufflenet_v2_x1_0(pretrained=True) model.fc = nn.Linear(model.fc.in_features, num_classes) # model.fc.out_features = num_classes optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9) elif model_name == 'ShuffleNet_v2_x1_5': model = models.shufflenet_v2_x1_5(pretrained=True) model.fc = nn.Linear(model.fc.in_features, num_classes) # model.fc.out_features = num_classes optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9) elif model_name == 'SqueezeNet 1_0': model = models.squeezenet1_0(pretrained=True) model.classifier[1] = nn.Conv2d(model.classifier[1].in_channels, num_classes,model.classifier[1].kernel_size, model.classifier[1].stride) # model.classifier[1].out_channels = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) elif model_name == 'SqueezeNet 1_1': model = models.squeezenet1_1(pretrained=True) model.classifier[1] = nn.Conv2d(model.classifier[1].in_channels, num_classes,model.classifier[1].kernel_size, model.classifier[1].stride) # model.classifier[1].out_channels = num_classes optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9) exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) criterion = nn.CrossEntropyLoss() model_ft = self.train(model, criterion, optimizer, exp_lr_scheduler, num_epochs=epochs) torch.save(model.state_dict(), 'model.pt') return model_ft def imshow(self,inp, title=None): """Display image for Tensor.""" inp = inp.numpy().transpose((1, 2, 0)) mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) inp = std * inp + mean inp = np.clip(inp, 0, 1) plt.imshow(inp) if title is not None: plt.title(title) plt.pause(0.001) def visualize_model(self,model, num_images=6): was_training = model.training model.eval() images_so_far = 0 fig = plt.figure() with torch.no_grad(): for i, (inputs, labels) in enumerate(self.dataloaders['val']): inputs = inputs.to(self.device) labels = labels.to(self.device) outputs = model(inputs) _, preds = torch.max(outputs, 1) for j in range(inputs.size()[0]): images_so_far += 1 ax = plt.subplot(num_images//2, 2, images_so_far) ax.axis('off') ax.set_title(f'predicted: {self.class_names[preds[j]]}') self.imshow(inputs.cpu().data[j]) if images_so_far == num_images: model.train(mode=was_training) return model.train(mode=was_training) def pytorch_predict(self,model): ''' Make prediction from a pytorch model ''' # set model to evaluate model model.eval() y_true = torch.tensor([], dtype=torch.long, device=self.device) all_outputs = torch.tensor([], device=self.device) # deactivate autograd engine and reduce memory usage and speed up computations with torch.no_grad(): for data in self.dataloaders['test']: inputs = [i.to(self.device) for i in data[:-1]] labels = data[-1].to(self.device) outputs = model(*inputs) y_true = torch.cat((y_true, labels), 0) all_outputs = torch.cat((all_outputs, outputs), 0) y_true = y_true.cpu().numpy() _, y_pred = torch.max(all_outputs, 1) y_pred = y_pred.cpu().numpy() y_pred_prob = F.softmax(all_outputs, dim=1).cpu().numpy() return y_true, y_pred, y_pred_prob