File size: 2,591 Bytes
b39a019 |
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 67 68 69 70 71 72 73 74 |
import torch
import CrosswalkDataset as Dataset
import ClassifierModel as Model
import Utilities as Utils
def train_model_v0(model_to_train, dataset, epoch_number=25, loss_func=Utils.BasicClassificationLoss,
batch_size=16, save=False):
optimiser = torch.optim.Adam(model_to_train.parameters(), lr=0.001)
dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=batch_size)
loss_function = loss_func()
for epoch in range(epoch_number):
model_to_train.train()
running_loss = 0.0
for images, gt_labels in dataloader:
optimiser.zero_grad()
predictions = model_to_train(images)
batch_loss = loss_function(predictions, gt_labels)
batch_loss.backward()
running_loss += batch_loss
optimiser.step()
print(f"Epoch [{epoch + 1} of {epoch_number}] finished, with loss {running_loss / len(dataloader)} in "
f"len {len(dataloader) * batch_size}")
Utils.save_model(model_to_train, optimiser)
return model_to_train
# Additionally incorporated a learning rate scheduler
def train_model_v1(model_to_train, dataset, epoch_number=10, loss_func=Utils.BasicClassificationLoss,
batch_size=16, save=False):
optimiser = torch.optim.Adam(model_to_train.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimiser, gamma=0.95)
dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=batch_size)
loss_function = loss_func()
for epoch in range(epoch_number):
model_to_train.train()
running_loss = 0.0
for images, gt_labels in dataloader:
optimiser.zero_grad()
predictions = model_to_train(images)
batch_loss = loss_function(predictions, gt_labels)
batch_loss.backward()
running_loss += batch_loss
optimiser.step()
scheduler.step()
print(f"Epoch [{epoch + 1} of {epoch_number}] finished, with loss {running_loss / len(dataloader)} in "
f"len {len(dataloader) * batch_size}")
Utils.save_model(model_to_train, optimiser)
return model_to_train
model = Model.BasicClassificationModel(image_size=416)
# size should be dynamically obtained later on
crosswalk_dataset = Dataset.CrosswalkDataset("Crosswalk.v7-crosswalk-t3.tensorflow/train/_annotations.csv",
"Crosswalk.v7-crosswalk-t3.tensorflow/train")
model = train_model_v1(model, crosswalk_dataset, save=True)
|