from engine import train_one_epoch, evaluate

# train on the accelerator or on the CPU, if an accelerator is not available
[device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device") = [torch.accelerator.current_accelerator](https://docs.pytorch.org/docs/stable/generated/torch.accelerator.current_accelerator.html#torch.accelerator.current_accelerator "torch.accelerator.current_accelerator")() if [torch.accelerator.is_available](https://docs.pytorch.org/docs/stable/generated/torch.accelerator.is_available.html#torch.accelerator.is_available "torch.accelerator.is_available")() else [torch.device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device")('cpu')

# our dataset has two classes only - background and person
num_classes = 2
# use our dataset and defined transformations
[dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset") = [PennFudanDataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "torch.utils.data.Dataset")('data/PennFudanPed', get_transform(train=True))
[dataset_test](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset") = [PennFudanDataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "torch.utils.data.Dataset")('data/PennFudanPed', get_transform(train=False))

# split the dataset in train and test set
indices = [torch.randperm](https://docs.pytorch.org/docs/stable/generated/torch.randperm.html#torch.randperm "torch.randperm")(len([dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset"))).tolist()
[dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset") = [torch.utils.data.Subset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset")([dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset"), indices[:-50])
[dataset_test](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset") = [torch.utils.data.Subset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset")([dataset_test](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset"), indices[-50:])

# define training and validation data loaders
[data_loader](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader") = [torch.utils.data.DataLoader](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader")(
    [dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset"),
    batch_size=2,
    shuffle=True,
    collate_fn=utils.collate_fn
)

[data_loader_test](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader") = [torch.utils.data.DataLoader](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader")(
    [dataset_test](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Subset "torch.utils.data.Subset"),
    batch_size=1,
    shuffle=False,
    collate_fn=utils.collate_fn
)

# get the model using our helper function
model = get_model_instance_segmentation(num_classes)

# move model to the right device
[model.to](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.to "torch.nn.Module.to")([device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device"))

# construct an optimizer
params = [p for p in [model.parameters](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.parameters "torch.nn.Module.parameters")() if p.requires_grad]
[optimizer](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD "torch.optim.SGD") = [torch.optim.SGD](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD "torch.optim.SGD")(
    params,
    lr=0.005,
    momentum=0.9,
    weight_decay=0.0005
)

# and a learning rate scheduler
[lr_scheduler](https://docs.pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html#torch.optim.lr_scheduler.StepLR "torch.optim.lr_scheduler.StepLR") = [torch.optim.lr_scheduler.StepLR](https://docs.pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html#torch.optim.lr_scheduler.StepLR "torch.optim.lr_scheduler.StepLR")(
    [optimizer](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD "torch.optim.SGD"),
    step_size=3,
    gamma=0.1
)

# let's train it just for 2 epochs
num_epochs = 2

for epoch in range(num_epochs):
    # train for one epoch, printing every 10 iterations
    train_one_epoch(model, [optimizer](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD "torch.optim.SGD"), [data_loader](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader"), [device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device"), epoch, print_freq=10)
    # update the learning rate
    [lr_scheduler.step](https://docs.pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html#torch.optim.lr_scheduler.StepLR.step "torch.optim.lr_scheduler.StepLR.step")()
    # evaluate on the test dataset
    evaluate(model, [data_loader_test](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader"), [device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device")=[device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device"))

print("That's it!")