CoverageControl/python/training/train_lpac.py at main · KumarRobotics/CoverageControl · GitHub

History

131 lines (105 loc) · 4.3 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

"""

Train the LPAC model

"""

# @file train_lpac.py

# @brief Train the LPAC model

import os

import pathlib

import sys

import torch

import torch_geometric

from coverage_control import IOUtils

from coverage_control.nn import CNNGNNDataset

from coverage_control.nn import LPAC

from coverage_control.nn import TrainModel

def main():

# Set the device

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

config_file = sys.argv[1]

world_size = int(sys.argv[2])

config = IOUtils.load_toml(config_file)

num_workers = config["NumWorkers"]

dataset_path = pathlib.Path(IOUtils.sanitize_path(config["DataDir"]))

data_dir = dataset_path / "data/"

lpac_model = config["LPACModel"]

model_dir = IOUtils.sanitize_path(lpac_model["Dir"]) + "/"

if not os.path.exists(model_dir):

os.makedirs(model_dir)

training_config = config["LPACTraining"]

batch_size = training_config["BatchSize"]

num_epochs = training_config["NumEpochs"]

learning_rate = training_config["LearningRate"]

# momentum = training_config["Momentum"]

weight_decay = training_config["WeightDecay"]

use_comm_map = config["ModelConfig"]["UseCommMaps"]

model = LPAC(config).to(device)

# model = torch.compile(model, dynamic=True)

# cnn_pretrained_model = config["CNNModel"]["Dir"] + config["CNNModel"]["Model"]

# model.LoadCNNBackBone(cnn_pretrained_model)

if "PreTrainedModel" in config["LPACModel"]:

lpac_pretrained_model = (

IOUtils.sanitize_path(config["LPACModel"]["Dir"])

+ "/" + config["LPACModel"]["PreTrainedModel"]

)

model.load_model(lpac_pretrained_model)

target_type = "actions"

if "TargetType" in config["ModelConfig"]:

target_type = config["ModelConfig"]["TargetType"]

train_dataset = CNNGNNDataset(data_dir, "train", use_comm_map, world_size, target_type)

val_dataset = CNNGNNDataset(data_dir, "val", use_comm_map, world_size, target_type)

# Check if buffer exists

if not hasattr(model, "actions_mean"):

model.register_buffer("actions_mean", train_dataset.targets_mean.to(device))

model.register_buffer("actions_std", train_dataset.targets_std.to(device))

else:

model.actions_mean = train_dataset.targets_mean

model.actions_std = train_dataset.targets_std

print("Loaded datasets")

print(f"Train dataset size: {len(train_dataset)}")

# Python 3.14 defaults to the forkserver start method on Linux, which

# re-imports __main__ and hands the dataset tensors to workers through

# shared memory. Fork keeps copy-on-write access to the datasets already

# loaded in this process. Workers only collate CPU tensors, so forking

# after CUDA init is safe here.

mp_context = "fork" if num_workers > 0 else None

train_loader = torch_geometric.loader.DataLoader(

train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers,

multiprocessing_context=mp_context

)

val_loader = torch_geometric.loader.DataLoader(

val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers,

multiprocessing_context=mp_context

)

# optimizer = torch.optim.SGD(

# model.parameters(),

# lr=learning_rate,

# momentum=momentum,

# weight_decay=weight_decay

# )

optimizer = torch.optim.Adam(

model.parameters(),

lr=learning_rate,

weight_decay=weight_decay

)

# Use mse loss for regression

criterion = torch.nn.MSELoss()

trainer = TrainModel(

model,

train_loader,

val_loader,

optimizer,

criterion,

num_epochs,

device,

model_dir,

)

# trainer = TrainModel(model, train_loader, val_loader, optimizer, criterion, num_epochs, device, model_dir)

trainer.train()

test_dataset = CNNGNNDataset(data_dir, "test", use_comm_map, world_size, target_type)

test_loader = torch_geometric.loader.DataLoader(

test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers,

multiprocessing_context=mp_context

)

test_loss = trainer.test(test_loader)

if __name__ == "__main__":

main()

CoverageControl/python/training/train_lpac.py at main · KumarRobotics/CoverageControl · GitHub

Uh oh!

Expand file tree

Breadcrumbs

train_lpac.py

Latest commit

History

Breadcrumbs

train_lpac.py

File metadata and controls