def _update(engine, batch): x, indices = _prepare_batch(batch) y_pred = model(x) y_pred = F.softmax(y_pred, dim=1) return { "y_pred": convert_tensor(y_pred, device='cpu'), "indices": indices }
def __call__(self, img, faces): if len(faces) == 0: return [] # crop, face, target = self.transform((img, [face], img.size, target)) # return ((*crop, *face), target) t = [self.transform((img, [f], img.size)) for f in faces] # t = self.transform((img, faces, img.size)) t = [(*x[0], *x[1]) for x in t] batch = default_collate(t) batch = convert_tensor(batch, device='cuda') preds = self.model(batch).data.cpu() gaze_targets = [ cam2screen(p[0], p[1], self.device_params) for p in preds ] return gaze_targets
def _prepare_batch(batch): (x, y), indices = batch x = convert_tensor(x, device=device) return x, y, indices