Python 重新训练 pytorch 模型(增强学习)

我有以下代码:

import torch
from facenet_pytorch import InceptionResnetV1, MTCNN
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
import pandas as pd
import os


workers = 0 if os.name == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))

mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
    device=device
)

def collate_fn(x):
    return x[0]

dataset = datasets.ImageFolder('data/images/')
dataset.idx_to_class = {i:c for c, i in dataset.class_to_idx.items()}
loader = DataLoader(dataset, collate_fn=collate_fn, num_workers=workers)
#print(dataset.idx_to_class)

aligned = []
names = []
i = 0
for x, y in loader:
    x_aligned, prob = mtcnn(x, return_prob=True)
    if x_aligned is not None:
        print('Face detected with probability: {:8f}'.format(prob))
        aligned.append(x_aligned)
        names.append(dataset.idx_to_class[y])
        i += 1
#print(i)

for name, param in mtcnn.named_parameters(): #Freezing everything but last layer
    #print(name)
    if name != "onet.dense6_3.bias":
        param.require_grad = False
    else:
        param.require_grad = True

现在我想重新训练这个模型来预测三个等级(现在它只预测一张脸的概率).假设我有data/images/个文件夹，分别是faces1、faces2和faces3.如何使用这三个文件夹重新训练此模型？我想要一个像[prob1, prob2, prob3]这样的张量，每个类都有一个图像的概率.谢谢