File size: 1,765 Bytes
1c13d92 |
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 |
import torch
import torch.nn as nn
import torch.optim as optim
from definition import FlowersImagesDetectionModel
from torch.utils.data import DataLoader
from datasets import load_dataset
from torchvision.transforms import ToTensor, Resize
from torch.utils.data.dataset import TensorDataset
flowerTypesNumber = 102
model = FlowersImagesDetectionModel(flowerTypesNumber)
# Funzioni di ottimizzazione e di perdita
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# Caricamento del dataset
originalDataset = load_dataset("nelorth/oxford-flowers", split="train")
tensorImages = []
tensorLabels = []
# Trasforma le immagini in tensori PyTorch e ridimensionale
for imageData, label in zip(originalDataset['image'], originalDataset['label']):
tensorImage = ToTensor()(Resize((224, 224))(imageData)) # Ridimensiona le immagini
tensorImages.append(tensorImage)
tensorLabels.append(label)
# Trasforma le liste di tensori in un singolo tensore
imagesTensor = torch.stack(tensorImages)
labelsTensor = torch.tensor(tensorLabels)
# Crea un dataset
dataset = TensorDataset(imagesTensor, labelsTensor)
# Crea un DataLoader
dataLoader = DataLoader(dataset, batch_size=64, shuffle=True)
# Addestramento
model.train()
for epoch in range(2):
running_loss = 0.0
for i, (inputs, labels) in enumerate(dataLoader, 0):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 99:
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 100))
running_loss = 0.0
torch.save(model.state_dict(), 'pytorch_model.bin')
|