File size: 2,532 Bytes
f373456 cdb9930 f373456 cdb9930 f373456 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
# import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision.models import resnet50, ResNet50_Weights
from utils import *
from data_setup import classes
class ImageClassificationBase(nn.Module):
def training_step(self, batch):
images, labels = batch
out = self(images)
# labels = labels.float().unsqueeze(1)
loss = F.cross_entropy(out, labels)
acc = accuracy(out, labels)
# print('training loss and acc:', loss, acc)
return loss, acc
def validation_step(self, batch):
images, labels = batch
out = self(images)
# labels = labels.float().unsqueeze(1)
loss = F.cross_entropy(out, labels)
acc = accuracy(out, labels)
# print('Validation loss and acc:', loss, acc)
return {'val_loss':loss.detach(), 'val_acc':acc}
def validation_end_epoch(self, results):
batch_loss = [x['val_loss'] for x in results]
epoch_loss = torch.stack(batch_loss).mean()
batch_acc = [x['val_acc'] for x in results]
epoch_acc = torch.stack(batch_acc).mean()
return {'val_loss':epoch_loss.item(), 'val_acc':epoch_acc.item()}
# def epoch_end(self, epoch, outputs):
# print(f"Epoch {epoch+1}: train_loss: {outputs['train_loss']}, val_loss: {outputs['val_loss']}, val_acc: {outputs['val_acc']}")
def epoch_end(self, epoch, result):
print(f"Epoch {epoch+1}: train_loss: {result['train_losses']:.4f}, train_acc: {result['train_acc']:.4f}, \
val_loss: {result['val_loss']:.4f}, val_acc: {result['val_acc']:.4f} ")
class FlowerClassificationModel(ImageClassificationBase):
def __init__(self, num_classes, pretrained=True):
super().__init__()
if pretrained:
if torchvision.__version__ >= '0.13.0':
self.network = self.network = resnet50(weights=ResNet50_Weights.DEFAULT)
else:
# 1. Get the base model with pretrained weights and send to target device
self.network = torchvision.models.resnet50(pretrained=True)
for param in self.network.parameters():
param.requires_grad = False
# Replace last layer
self.network.fc = nn.Linear(in_features=2048,
out_features=num_classes, # same number of output units as our number of classes
bias=True)
def forward(self, xb):
return self.network(xb) |