/img/star (2).png.webp)
/img/news (2).png.webp)
/img/event.png.webp)
Время прочтения: 6 мин.
PyTorch — это отличная библиотека для решения разнообразных задач машинного обучения. Она сочетает в себе простоту и эффективность, поэтому является достаточно популярной. При этом, авторам было замечено, что во многих туториалах и руководствах написанный torch-код хоть и выполняет поставленную задачу, делает это не самым эффективным образом. Учитывая то, что люди крайне редко самостоятельно читают официальную документацию, будет полезным привести примеры некоторых простых способов оптимизации torch-кода.
Данная статья рассчитана на новичков в PyTorch, но возможно, некоторые советы помогут и тем, кто давно использует данную библиотеку.
Для начала рассмотрим простую, но крайне важную вещь: загрузку обучающей выборки в модель. Не секрет, что для этого используется модуль torch.unils.data.DataLoader. Пусть дан некий абстрактный класс датасета, создающий обучающую выборку из изображений и csv разметки:
1. def parse_one_annot(path_to_data_file, filename):
2. data = pd.read_csv(path_to_data_file)
3. boxes_array = data[data["filename"] == filename][["xmin", "ymin",
4. "xmax", "ymax"]].values
5. return boxes_array
6.
7.
8. def get_transform(train):
9. transforms = []
10. transforms.append(T.ToTensor())
11. if train:
12. transforms.append(T.RandomHorizontalFlip(0.5))
13. return T.Compose(transforms)
14.
15.
16. class AbstractDataset(torch.utils.data.Dataset):
17. def __init__(self, root, data_file, transforms):
18. self.root = root
19. self.transforms = transforms
20. self.imgs = sorted(os.listdir(os.path.join(root, "images")))
21. self.path_to_data_file = data_file
22.
23. def __getitem__(self, idx):
24. img_path = os.path.join(self.root, "images", self.imgs[idx])
25. img = Image.open(img_path).convert("RGB")
26. box_list = parse_one_annot(self.path_to_data_file, self.imgs[idx])
27. boxes = torch.as_tensor(box_list, dtype=torch.float32)
28. num_objs = len(box_list)
29. labels = torch.ones((num_objs,), dtype=torch.int64)
30. image_id = torch.tensor([idx])
31. area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
32. iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
33. target = {}
34. target["boxes"] = boxes
35. target["labels"] = labels
36. target["image_id"] = image_id
37. target["area"] = area
38. target["iscrowd"] = iscrowd
39. if self.transforms is not None:
40. img, target = self.transforms(img, target)
41. return img, target
42. def __len__(self):
43. return len(self.imgs)
Создадим нашу выборку:
1. data_loader = torch.utils.data.DataLoader(
2. dataset, batch_size=1, shuffle=True, num_workers=4,
3. collate_fn=utils.collate_fn, pin_memory=True)
Обратите внимание, на выделенный красный параметр num_workers. По умолчанию он равен нулю, и данный параметр отвечает за возможность количество потоков параллельно загрузки данных в модель. Конечно, на это тратится память вашей GPU, и памяти остается меньше на само обучение модели, однако с мощными видеокартами (cuda capability > 5.2) при значении параметра по умолчанию бывает так, что итерация на загруженном батче уже прошли, а новых данных еще не поступило, что при бльших датасетах может сильно учеличить длительность обучения.
Возникает вопрос, как определить оптимальное количество воркеров для ваше системы. Здесь существует универсальное правило: num_workers = 4 * num_GPU. Однако, если вы вдруг сталкиваетесь с memory allocation error, то конечно необходимо уменьшить множитель.
Следует заметить, что если вы работаете в Windows, то для использования данного параметра необходимо обязательно прописывать функцию main(), и в ней вызвать freeze_support(), иначе вы будете получать ошибку BrokenPipeError
1. from torch.multiprocessing import freeze_support
2. if __name__ == '__main__':
3. freeze_support()
К сожалению, данное решение работает также не у всех, и если вам не повезло, то либо придется использовать num_workers = 0, либо менять систему.
Также, будет полезным упомянуть, что следует использовать автоматическое выделение памяти GPU под ваши данные, устанавливая значения параметра pin_memory = True. Данное действие ограничивает возможность модели использовать всю доступную память видеокарты, которое на самом деле может быть избыточно. Кроме того, с этим параметром под данные будет выделяться столько памяти, сколько возможно выделить, не мешая обучению модели. Также, использование данного параметра освобождает нас от вызова небезопасно функции очистки кеша: torch.cuba.empty_cache().
Мы немного оптимизировали нашу загрузку, теперь рассмотрим простой, но, пожалуй, самый главный принцип оптимизации в PyTorch, при использовании GPU: Вы должны любой ценой избегать перемещения данных из CPU в GPU и обратно.
При таком перемещении, данные необходимо каждый раз адаптировать под устроство, а это значительно ухудшает быстродействие.
Характерным признаком того, что вы много раз перегоняете данные с GPU на CPU служит использование таких функций как:
- .cpu()
- .item()
- .numpy
По возможности, избегайте их вызова, используя .detach() вместо них.
Еще никогда не забывайте про использование GPU не только для обучения, но и для общей работы сети. В подавляющем количестве обучающих материалов, модель в режиме .eval() не отправляется на устротво. Суть в том, что для использования модели в таком режиме на GPU, нам также необходимо портировать данные на GPU, но у нас уже не используется удобный DataLoader.
Однако, на самом деле, портировать данные на GPU очень просто, нужно только использовать встроенный метод преобразования данных в класс torch.tensor.
1. from torchvision.transforms import ToTensor as tens
2. import cv2
3.
4.
5. def get_model(num_classes):
6. # load an object detection model pre-trained on COCO
7. model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
8. in_features = model.roi_heads.box_predictor.cls_score.in_features
9. model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
10. return model
11.
12.
13. device = torch.device("cuda")
14. print(torch.cuda.get_device_capability(device))
15. loaded_model = get_model(num_classes=2)
16. loaded_model.load_state_dict(torch.load('model.pth'))
17. loaded_model.to(device)
18. frame = cv2.imread('frame.jpg')
19. img_tens = tens()(frame)
20. cuda_tens = img_tens.cuda()
21. cuda_tens.to(device)
22. with torch.no_grad():
23. prediction = loaded_model([cuda_tens])
Использование данного метода в коде отмечено красным. Также, не забывайте использовать torch.no_grad() в режиме .eval() — это значительно снижает потребляемую память, так как .no_grad() отключает используемые слои в нейросети (dropout, batch_normalization) при ее непосредственном использовании.
Еще один способ значительно ускорить ваш torch-код это создать нужные вам объекты напрямую на GPU. Как мы помним, трансфер данных с CPU на GPU отнимает и время, и производительность, поэтому, если вы имеете необходимость в создании тензоров, делайте это не так:
1. t = tensor.rand(2,2).cuda()
а так:
2. t = tensor.rand(2,2, device=torch.device('cuda:0'))
В первом случае, вы создаете тензор на CPU, а затем трансферите его на GPU (также мы делаем в предыдущем примере, но в том случае, мы работаем с заданными изображениями, и не можем их считывать в память видеокарты (на самом деле, наверное, можем, но я пока не знаю как))
К сожалению, не всегда возможно создание таких тензоров, но если, например, вы работаете с генеративно-состязательными сетями, выход генератора можно сделать напрямую в видеопамять.
В качестве заключения, из всей статьи можно выделить, что суть оптимизации torch-кода сводится к простым правилам:
- используйте данные вам встроенные возможности оптимизации, такие как распараллеливание загрузки данных на устройство, выделение требуемого количества памяти под него, отключение “лишних” слоев сети в режиме .eval() и прочее,
- Как можно меньше переносите данные из CPU на GPU и обратно,
- Старайтесь при возможности сразу создавать данные на GPU.
На самом деле, существует еще множество более продвинутых способов оптимизации кода на PyTorch, но даже если вы будете следовать данным простым правилам, вы сэкономите немало вычислительного времени.