HRNetベースのライブラリを自前データセットを用いて実行したい

前提

以下のライブラリを用いてランドマークを検出するモデルを学習しようと考えています。
https://github.com/HRNet/HRNet-Facial-Landmark-Detection

こちらの記事を参考に300Wデータセットを用いたモデルの学習・テストを実行し、正しく動作したことを確認できました。
https://tadaoyamaoka.hatenablog.com/entry/2022/06/12/153543

実現したいこと

300Wデータセットを自前のデータセットに置き換えてモデルの学習を行いたいと考えています。
300Wデータセットでは68点のランドマークを学習するのに対し、自前データセットでは22点のランドマークを学習に使います。

発生している問題・エラーメッセージ

300Wデータセットでモデルの学習を行ったのと同じように、tools/train.pyを実行したところ以下のようなエラーが発生しました。

C:\Users\***\anaconda3\envs\landmark\lib\site-packages\torch\nn\modules\loss.py:520: UserWarning: Using a target size (torch.Size([16, 22, 64, 64])) that is different to the input size (torch.Size([16, 68, 64, 64])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size. return F.mse_loss(input, target, reduction=self.reduction) Traceback (most recent call last): File "D:\***\HRNet\tools\train.py", line 145, in <module> main() File "D:\***\HRNet\tools\train.py", line 117, in main function.train(config, train_loader, model, criterion, File "D:\***\HRNet\tools\..\lib\core\function.py", line 66, in train loss = critertion(output, target) File "C:\Users\***\anaconda3\envs\landmark\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\***\anaconda3\envs\landmark\lib\site-packages\torch\nn\modules\loss.py", line 520, in forward return F.mse_loss(input, target, reduction=self.reduction) File "C:\Users\***\anaconda3\envs\landmark\lib\site-packages\torch\nn\functional.py", line 3111, in mse_loss expanded_input, expanded_target = torch.broadcast_tensors(input, target) File "C:\Users\***\anaconda3\envs\landmark\lib\site-packages\torch\functional.py", line 72, in broadcast_tensors return _VF.broadcast_tensors(tensors) # type: ignore[attr-defined] RuntimeError: The size of tensor a (68) must match the size of tensor b (22) at non-singleton dimension 1

ランドマーク点の数が300Wデータセットと異なるために発生したエラーであることは分かりましたが、
どこを修正すればよいのか分からず質問させていただきました。

試したこと

HRNet/lib/core/evaluation.py の以下の部分を修正しましたが解決できませんでした。

for i in range(N): pts_pred, pts_gt = preds[i, ], target[i, ] if L == 19: # aflw interocular = meta['box_size'][i] elif L == 29: # cofw interocular = np.linalg.norm(pts_gt[8, ] - pts_gt[9, ]) elif L == 68: # 300 ← L == 68 から L == 22 に変更 # interocular interocular = np.linalg.norm(pts_gt[36, ] - pts_gt[45, ]) elif L == 98: interocular = np.linalg.norm(pts_gt[60, ] - pts_gt[72, ]) else: raise ValueError('Number of landmarks is wrong') rmse[i] = np.sum(np.linalg.norm(pts_pred - pts_gt, axis=1)) / (interocular * L)

補足情報(FW/ツールのバージョンなど)

実行環境
・Windows10
・Python 3.9.13
・PyTorch 1.10
・CUDA 11.3
・cuDNN 8.4.0

コメントを投稿

0 コメント