127 lines
3.5 KiB
Python
127 lines
3.5 KiB
Python
# %%
|
|
from engine import train_one_epoch, evaluate
|
|
from model import Model
|
|
from data import iNaturalistDataset
|
|
import torch
|
|
import os
|
|
import datetime as dt
|
|
import json
|
|
import utils
|
|
import pandas as pd
|
|
import sys
|
|
if not os.path.exists("models/"):
|
|
os.mkdir("models")
|
|
|
|
if torch.cuda.is_available():
|
|
device = torch.device("cuda")
|
|
else:
|
|
device = torch.device("cpu")
|
|
|
|
default_model_root = "models/" + dt.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
default_model_path = default_model_root + ".pth"
|
|
default_model_info = default_model_root + ".json"
|
|
default_state_path = default_model_root + ".oth"
|
|
|
|
|
|
#species_list = set(["Poecile atricapillus", "Archilochus colubris", "Icterus galbula"])
|
|
|
|
csv_path = '/home/thebears/Seafile/Designs/ML/inaturalist_models/species_occurence.csv'
|
|
df = pd.read_csv(csv_path)
|
|
species_list = set(list(df[df['count']>1000]['species']))
|
|
|
|
#model_type = "fasterrcnn_mobilenet_v3_large_fpn"
|
|
#batch_size = 16
|
|
|
|
model_type = 'fasterrcnn_resnet50_fpn'
|
|
batch_size = 8
|
|
|
|
|
|
|
|
def run(model_name = None, epoch_start = 0):
|
|
val_dataset = iNaturalistDataset(
|
|
validation=True,
|
|
species=species_list,
|
|
)
|
|
train_dataset = iNaturalistDataset(
|
|
train=True,
|
|
species=species_list,
|
|
)
|
|
|
|
|
|
if model_name is None:
|
|
fresh_start = True
|
|
model_info = default_model_info
|
|
model_path = default_model_path
|
|
state_path = default_state_path
|
|
else:
|
|
fresh_start = False
|
|
model_info = model_name.rstrip('.pth').rstrip('.json')+'.json'
|
|
model_path = model_info.rstrip('.json')+'.pth'
|
|
state_path = model_info.rstrip('.json')+'.oth'
|
|
|
|
|
|
if fresh_start:
|
|
with open(model_info, "w") as js_p:
|
|
json.dump(
|
|
{"categories": train_dataset.categories, "model_type": model_type},
|
|
js_p,
|
|
default=str,
|
|
indent=4,
|
|
)
|
|
|
|
train_data_loader = torch.utils.data.DataLoader(
|
|
train_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=10,
|
|
collate_fn=utils.collate_fn,
|
|
)
|
|
|
|
val_data_loader = torch.utils.data.DataLoader(
|
|
val_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=10,
|
|
collate_fn=utils.collate_fn,
|
|
)
|
|
|
|
num_classes = len(species_list) + 1
|
|
model = Model(num_classes, model_type)
|
|
model.to(device)
|
|
|
|
if not fresh_start:
|
|
model.load_state_dict(
|
|
torch.load(model_path, map_location = torch.device(device))
|
|
)
|
|
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
|
|
|
|
if os.path.exists(state_path):
|
|
optimizer.load_state_dict(torch.load(state_path, map_location = torch.device(device)))
|
|
|
|
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
|
|
|
|
num_epochs = 10
|
|
|
|
for epoch in range(epoch_start, num_epochs):
|
|
train_one_epoch(
|
|
model, optimizer, train_data_loader, device, epoch, print_freq=10 )
|
|
lr_scheduler.step()
|
|
torch.save(model.state_dict(), model_path)
|
|
torch.save(optimizer.state_dict(), state_path)
|
|
evaluate(model, val_data_loader, device=device)
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) == 3:
|
|
model_name = sys.argv[1]
|
|
epoch_start = int(sys.argv[2])
|
|
run(model_name = model_name, epoch_start = epoch_start)
|
|
else:
|
|
run()
|
|
|
|
# run()
|