Vit Artworkclassifier

🚀 vit-artworkclassifier

このモデルは、入力された任意の画像のアートワークスタイルを返します。

このモデルは、google/vit-base-patch16-224-in21k を imagefolder データセットでファインチューニングしたバージョンです。これは artbench - 10 データセット (https://www.kaggle.com/datasets/alexanderliao/artbench10) のサブセットで、各クラスに 1000 件のアートワークを含むトレーニングセットと、各クラスに 100 件のアートワークを含む検証セットがあります。 評価セットでは、以下の結果を達成しています。

• 損失: 1.1392
• 正解率: 0.5948

📚 ドキュメント

モデルの説明

このモデルがトレーニングされたプロジェクトの説明はこちらにあります: https://medium.com/@oliverpj.schamp/training-and-evaluating-stable-diffusion-for-artwork-generation-b099d1f5b7a6

想定される用途と制限

このモデルは、artbench - 10 の 10 クラスのうち 9 クラスのみを含んでおり、浮世絵 (ukiyo_e) は含まれていません。これは、データの入手可能性とフォーマットの問題によるものです。

トレーニングと評価データ

トレーニング: artbench - 10 からランダムに選択された 1000 枚の画像 (各クラス)。検証: artbench - 10 からランダムに選択された 100 枚の画像 (各クラス)。

トレーニング手順

トレーニングハイパーパラメータ

トレーニング中に使用されたハイパーパラメータは以下の通りです。

• 学習率: 0.0001
• トレーニングバッチサイズ: 32
• 評価バッチサイズ: 8
• シード: 42
• オプティマイザ: Adam (ベータ=(0.9, 0.999)、イプシロン=1e - 08)
• 学習率スケジューラのタイプ: 線形
• エポック数: 4
• 混合精度トレーニング: Native AMP

トレーニング結果

フレームワークのバージョン

• Transformers 4.26.1
• Pytorch 1.13.1+cu117
• Datasets 2.9.0
• Tokenizers 0.13.2

実行コード

def vit_classify(image):
    vit = ViTForImageClassification.from_pretrained("oschamp/vit-artworkclassifier")
    vit.eval()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    vit.to(device)
    
    model_name_or_path = 'google/vit-base-patch16-224-in21k'
    feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)

    #LOAD IMAGE

    encoding = feature_extractor(images=image, return_tensors="pt")
    encoding.keys()

    pixel_values = encoding['pixel_values'].to(device)

    outputs = vit(pixel_values)
    logits = outputs.logits

    prediction = logits.argmax(-1)
    return prediction.item() #vit.config.id2label[prediction.item()]

📄 ライセンス

このモデルは Apache - 2.0 ライセンスの下で提供されています。

📦 モデル情報