48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torchvision import models
|
|
import time
|
|
|
|
# Base model
|
|
from architectures.Model import Model
|
|
|
|
class Inception(Model):
|
|
def get(self):
|
|
m = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT)
|
|
#for param in model.parameters():
|
|
# param.requires_grad = False
|
|
m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, self.size)
|
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
|
return m
|
|
|
|
def train(self, epochs, loader, rate):
|
|
# Override because Inception returns a tuple (main, aux)
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)
|
|
|
|
print(f"Starting training on {self.device}...")
|
|
start_time = time.time()
|
|
|
|
self.model.train()
|
|
|
|
for epoch in range(epochs):
|
|
total_loss = 0.0
|
|
for inputs, labels in loader:
|
|
|
|
inputs, labels = inputs.to(self.device), labels.to(self.device)
|
|
optimizer.zero_grad()
|
|
|
|
outputs, aux_outputs = self.model(inputs)
|
|
loss = criterion(outputs, labels) + 0.3 * criterion(aux_outputs, labels)
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
total_loss += loss.item()
|
|
|
|
print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(loader):.4f}")
|
|
|
|
if self.device.type == 'cuda': torch.cuda.synchronize()
|
|
print(f"Training completed in: {time.time() - start_time:.2f}s")
|