def main(args):
input_size = 3
output_size = 1
# Make the training and testing set to be less bias
fire_dataset = FireDataset(args.csv_path[0])
fire_train, fire_test = random_split(
fire_dataset,
(round(0.7 * len(fire_dataset)), round(0.3 * len(fire_dataset))))
trainloader = DataLoader(fire_train,
batch_size=4096,
shuffle=True,
num_workers=2)
testloader = DataLoader(fire_test,
batch_size=512,
shuffle=False,
num_workers=2)
save_weights_pth = args.weights_path[0]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MLP(input_size=input_size, output_size=output_size)
model.to(device)
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=5e-7)
epochs = 30
if args.eval_only:
do_test(model, device, testloader, save_weights_pth)
else:
do_train(model, device, trainloader, criterion, optimizer, epochs,
save_weights_pth)
do_test(model, device, testloader, save_weights_pth)
def main(args):
input_size = 5
output_size = 3
epochs = 100
# Make the training and testing set to be less bias
hlb_dataset = HLBDataset(args.dataset_path[0])
hlb_train, hlb_test = random_split(
hlb_dataset,
(round(0.8 * len(hlb_dataset)), round(0.2 * len(hlb_dataset))))
print(args.weights_path[0])
print(f'Number of training examples: {len(hlb_train)}')
print(f'Number of testing examples: {len(hlb_test)}')
trainloader = DataLoader(hlb_train,
batch_size=2048,
shuffle=True,
num_workers=2)
testloader = DataLoader(hlb_test,
batch_size=1024,
shuffle=False,
num_workers=2)
save_weights_pth = args.weights_path[0]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MLP(input_size=input_size, output_size=output_size)
print(model)
model.to(device)
weights = torch.tensor([0.85, 0.25, 1.0])
criterion = nn.CrossEntropyLoss(weight=weights).cuda()
learning_rate = 1e-6
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
train_acc_array = []
train_loss_array = []
test_acc_array = []
test_loss_array = []
for i in range(epochs): # loop over the dataset multiple times
train_loss, train_acc = do_train(model, device, trainloader, criterion,
optimizer)
print('Epoch {} Train loss: {} Train acc: {}'.format(
i, train_loss, train_acc))
train_acc_array.append(train_acc)
train_loss_array.append(train_loss)
test_loss, test_acc = do_test(model, device, testloader, criterion)
test_acc_array.append(test_acc)
test_loss_array.append(test_loss)
print('Test loss: {} Test acc: {}'.format(test_loss, test_acc))
if i % 50 == 0:
learning_rate = learning_rate * 0.99
for param_group in optimizer.param_groups:
param_group['lr'] = learning_rate
torch.save(model.state_dict(), args.weights_path[0])
save_data(train_loss_array, train_acc_array, test_loss_array,
test_acc_array)
plot_graph(train_loss_array, train_acc_array, test_loss_array,
test_acc_array)