92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
from confic import (DEVICE, NUM_CLASSES, NUM_EPOCHS, OUTDIR, NUM_WORKERS, TRAIN_DIR)
|
|
from model import create_model
|
|
|
|
from tqdm.auto import tqdm
|
|
|
|
from datasets import create_train_test_dataset, create_train_loader, create_valid_loader
|
|
from custom_utils import Averager, SaveBestModel, save_model, save_loss_plot
|
|
|
|
import torch
|
|
import matplotlib.pyplot as plt
|
|
import time
|
|
|
|
from IPython import embed
|
|
|
|
def train(train_loader, model, optimizer):
|
|
print('Training')
|
|
global train_loss_list
|
|
|
|
prog_bar = tqdm(train_loader, total=len(train_loader))
|
|
for samples, targets in prog_bar:
|
|
images = list(image.to(DEVICE) for image in samples)
|
|
|
|
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]
|
|
|
|
loss_dict = model(images, targets)
|
|
|
|
losses = sum(loss for loss in loss_dict.values())
|
|
loss_value = losses.item()
|
|
train_loss_hist.send(loss_value) # this is a global instance !!!
|
|
train_loss_list.append(loss_value) # check what exactly this does !!!
|
|
|
|
optimizer.zero_grad()
|
|
losses.backward()
|
|
optimizer.step()
|
|
|
|
prog_bar.set_description(desc=f"Loss: {loss_value:.4f}")
|
|
|
|
return train_loss_list
|
|
|
|
|
|
if __name__ == '__main__':
|
|
train_data, test_data = create_train_test_dataset(TRAIN_DIR)
|
|
train_loader = create_train_loader(train_data)
|
|
test_loader = create_train_loader(test_data)
|
|
|
|
model = create_model(num_classes=NUM_CLASSES)
|
|
model = model.to(DEVICE)
|
|
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)
|
|
|
|
train_loss_hist = Averager()
|
|
val_loss_hist = Averager()
|
|
# train_itr = 1
|
|
# val_itr = 1
|
|
train_loss_list = []
|
|
val_loss_list = []
|
|
|
|
save_best_model = SaveBestModel()
|
|
|
|
for epoch in range(NUM_EPOCHS):
|
|
|
|
train_loss_hist.reset()
|
|
val_loss_hist.reset()
|
|
|
|
train_loss = train(train_loader, model, optimizer)
|
|
# val_loss = validate(train_loader, model, optimizer)
|
|
|
|
save_best_model(
|
|
val_loss_hist.value, epoch, model, optimizer
|
|
)
|
|
|
|
save_model(epoch, model, optimizer)
|
|
|
|
save_loss_plot(OUTDIR, train_loss, val_loss)
|
|
|
|
# prog_bar = tqdm(train_loader, total=len(train_loader))
|
|
# for samples, targets in prog_bar:
|
|
# images = list(image.to(DEVICE) for image in samples)
|
|
#
|
|
# targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]
|
|
#
|
|
# loss_dict = model(images, targets)
|
|
#
|
|
# losses = sum(loss for loss in loss_dict.values())
|
|
# loss_value = losses.item()
|
|
#
|
|
# optimizer.zero_grad()
|
|
# losses.backward()
|
|
# optimizer.step()
|
|
#
|
|
# prog_bar.set_description(desc=f"Loss: {loss_value:.4f}") |