Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn.functional as F | |
| import torch.nn as nn | |
| class Base(nn.Module): | |
| def training_step(self, batch): | |
| images, labels = batch | |
| out = self(images) # Generate predictions | |
| loss = F.cross_entropy(out, labels) # Calculate loss | |
| return loss | |
| def validation_step(self, batch): | |
| images, labels = batch | |
| out = self(images) # Generate predictions | |
| loss = F.cross_entropy(out, labels) # Calculate loss | |
| acc = accuracy(out, labels) # Calculate accuracy | |
| return {'val_loss': loss.detach(), 'val_acc': acc} | |
| def validation_epoch_end(self, outputs): | |
| batch_losses = [x['val_loss'] for x in outputs] | |
| epoch_loss = torch.stack(batch_losses).mean() # Combine losses | |
| batch_accs = [x['val_acc'] for x in outputs] | |
| epoch_acc = torch.stack(batch_accs).mean() # Combine accuracies | |
| return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()} | |
| def epoch_end(self, epoch, result): | |
| print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format( | |
| epoch, result['train_loss'], result['val_loss'], result['val_acc'])) | |
| # print(f'Epoch: {epoch} | Train_loss: {result['train_loss']} | Val_loss:{result['val_loss']} | Val_acc: {result['val_acc']}') | |
| def accuracy(outputs, labels): | |
| _, preds = torch.max(outputs, dim=1) | |
| return torch.tensor(torch.sum(preds == labels).item() / len(preds)) | |
| class PotatoDiseaseDetectionModel(Base): | |
| def __init__(self, in_channels=3, num_classes=3): | |
| super(PotatoDiseaseDetectionModel, self).__init__() | |
| # Define the network layers | |
| self.network = nn.Sequential( | |
| nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(kernel_size=2, stride=2), | |
| nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(kernel_size=2, stride=2), | |
| nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(kernel_size=2, stride=2), | |
| nn.Flatten() | |
| ) | |
| # Define the classifier layers | |
| self.classifier = nn.Sequential( | |
| nn.Linear(in_features=256*28*28, out_features=128), | |
| nn.BatchNorm1d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(0.5), | |
| nn.Linear(in_features=128, out_features=num_classes) | |
| ) | |
| def forward(self, x): | |
| # Pass the input through the network | |
| x = self.network(x) | |
| # Pass the output through the classifier | |
| x = self.classifier(x) | |
| return x | |
| # Create the model with desired number of classes | |
| potato_model = PotatoDiseaseDetectionModel(num_classes=3) | |
| tomato_model = PotatoDiseaseDetectionModel(num_classes=3) |