def __init__(self, context: det.TrialContext) -> None:
self.context = context
# Create a unique download directory for each rank so they don't
# overwrite each other.
self.download_directory = f"/tmp/data-rank{self.context.distributed.get_rank()}"
download_data(
download_directory=self.download_directory, data_config=self.context.get_data_config(),
)
dataset = PennFudanDataset(self.download_directory + "/PennFudanPed", get_transform())
# Split 80/20 into training and validation datasets.
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
self.dataset_train, self.dataset_val = torch.utils.data.random_split(
dataset, [train_size, test_size]
)
def __init__(self, context: PyTorchTrialContext) -> None:
self.context = context
# Create a unique download directory for each rank so they don't
# overwrite each other.
self.download_directory = f"/tmp/data-rank{self.context.distributed.get_rank()}"
download_data(
download_directory=self.download_directory, data_config=self.context.get_data_config(),
)
dataset = PennFudanDataset(self.download_directory + "/PennFudanPed", get_transform())
# Split 80/20 into training and validation datasets.
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
self.dataset_train, self.dataset_val = torch.utils.data.random_split(
dataset, [train_size, test_size]
)
model = fasterrcnn_resnet50_fpn(pretrained=True)
# Replace the classifier with a new two-class classifier. There are
# only two "classes": pedestrian and background.
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# Wrap the model.
self.model = self.context.wrap_model(model)
# Wrap the optimizer.
self.optimizer = self.context.wrap_optimizer(torch.optim.SGD(
self.model.parameters(),
lr=self.context.get_hparam("learning_rate"),
momentum=self.context.get_hparam("momentum"),
weight_decay=self.context.get_hparam("weight_decay"),
))
# Wrap the LR scheduler.
self.lr_scheduler = self.context.wrap_lr_scheduler(
torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=3, gamma=0.1),
step_mode=LRScheduler.StepMode.STEP_EVERY_EPOCH
)
def main():
print("Torch version:", torch.__version__)
# get command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--model_name',
type=str,
default="pytorch-peds.pt",
help='name with which to register your model')
parser.add_argument('--output_dir',
default="local-outputs",
type=str,
help='output directory')
parser.add_argument('--n_epochs',
type=int,
default=10,
help='number of epochs')
args = parser.parse_args()
# In case user inputs a nested output directory
os.makedirs(name=args.output_dir, exist_ok=True)
# Get a dataset by name
root_dir = download_data()
# use our dataset and defined transformations
dataset = PennFudanDataset(root=root_dir,
transforms=get_transform(train=True))
dataset_test = PennFudanDataset(root=root_dir,
transforms=get_transform(train=False))
# split the dataset in train and test set
torch.manual_seed(1)
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=True,
num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn)
if torch.cuda.is_available():
print('Using GPU')
device = torch.device('cuda')
else:
print('Using CPU')
device = torch.device('cpu')
# our dataset has two classes only - background and person
num_classes = NUM_CLASSES
# get the model using our helper function
model = get_instance_segmentation_model(num_classes)
# move model to the right device
model.to(device)
# construct an optimizer
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)
# and a learning rate scheduler which decreases the learning rate by
# 10x every 3 epochs
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
for epoch in range(args.n_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
data_loader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)
# Saving the state dict is recommended method, per
# https://pytorch.org/tutorials/beginner/saving_loading_models.html
torch.save(model.state_dict(),
os.path.join(args.output_dir, args.model_name))
import transforms as T
from engine import train_one_epoch, evaluate
from model import get_model_instance_segmentation
from data import PennFudanDataset
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
trainset = PennFudanDataset('PennFudanPed', get_transform(train=True))
testset = PennFudanDataset('PennFudanPed', get_transform(train=False))
indeces = torch.randperm(len(trainset)).tolist()
trainset = torch.utils.data.Subset(trainset, indeces[:-50])
testset = torch.utils.data.Subset(testset, indeces[-50:])
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=2,
num_workers=4,
shuffle=True,
collate_fn=utils.collate_fn)
testloader = torch.utils.data.DataLoader(testset,
batch_size=1,
num_workers=4,
shuffle=True,