<gh_stars>0
import torch
import torch.nn as nn
import torchvision.datasets as datasets

from utils.worker_template import BasicWorker
from utils.meters import AverageMeter
from .loader import TwoCropsTransform
from .builder import PLMoCoPlus, PatchLevelEncoder


class Worker(BasicWorker):
    def _create_model(self):
        model = PLMoCoPlus(
            PatchLevelEncoder, self.args.pretrained,
            self.args.moco_dim, self.args.moco_k, self.args.moco_m, self.args.moco_t)
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
        print("=> model\n" + str(model))
        return model

    def _create_dataset(self):
        data_root = self.args.data_root
        transform_func = TwoCropsTransform(crop_size=self.args.crop_size)
        train_dataset = datasets.ImageFolder(data_root, transform=transform_func)
        print("=> dataset\n" + str(train_dataset))
        return train_dataset

    def _init_additional_meters(self):
        loss_moco = AverageMeter('Loss(MoCo)', ':.4e')
        acc_moco = AverageMeter('Acc(MoCo)', ':6.2f')
        loss_plmoco = AverageMeter('Loss(PLMoCo)', ':.4e')
        acc_plmoco = AverageMeter('Acc(PLMoCo)', ':6.2f')
        return [loss_moco, acc_moco, loss_plmoco, acc_plmoco]

    def _forward(self, data):
        # unpack data
        (images, mask_q), _ = data
        if self.args.gpu is not None:
            images[0] = images[0].cuda(self.args.gpu, non_blocking=True)
            images[1] = images[1].cuda(self.args.gpu, non_blocking=True)
            mask_q = mask_q.cuda(self.args.gpu, non_blocking=True)
        # compute output
        with torch.cuda.amp.autocast(enabled=self.args.fp16):
            loss, loss_moco, acc_moco, loss_plmoco, acc_plmoco = self.model(im_q=images[0], im_k=images[1], mask_q=mask_q)
        # return loss and additional info (as list)
        return loss, [loss_moco, acc_moco, loss_plmoco, acc_plmoco]
import torch
import torch.nn as nn
import torchvision.datasets as datasets

from utils.worker_template import BasicWorker
from utils.meters import AverageMeter
from .loader import TwoCropsTransform
from .builder import PLMoCoPlus, PatchLevelEncoder


class Worker(BasicWorker):
    def _create_model(self):
        model = PLMoCoPlus(
            PatchLevelEncoder, self.args.pretrained,
            self.args.moco_dim, self.args.moco_k, self.args.moco_m, self.args.moco_t)
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
        print("=> model\n" + str(model))
        return model

    def _create_dataset(self):
        data_root = self.args.data_root
        transform_func = TwoCropsTransform(crop_size=self.args.crop_size)
        train_dataset = datasets.ImageFolder(data_root, transform=transform_func)
        print("=> dataset\n" + str(train_dataset))
        return train_dataset

    def _init_additional_meters(self):
        loss_moco = AverageMeter('Loss(MoCo)', ':.4e')
        acc_moco = AverageMeter('Acc(MoCo)', ':6.2f')
        loss_plmoco = AverageMeter('Loss(PLMoCo)', ':.4e')
        acc_plmoco = AverageMeter('Acc(PLMoCo)', ':6.2f')
        return [loss_moco, acc_moco, loss_plmoco, acc_plmoco]

    def _forward(self, data):
        # unpack data
        (images, mask_q), _ = data
        if self.args.gpu is not None:
            images[0] = images[0].cuda(self.args.gpu, non_blocking=True)
            images[1] = images[1].cuda(self.args.gpu, non_blocking=True)
            mask_q = mask_q.cuda(self.args.gpu, non_blocking=True)
        # compute output
        with torch.cuda.amp.autocast(enabled=self.args.fp16):
            loss, loss_moco, acc_moco, loss_plmoco, acc_plmoco = self.model(im_q=images[0], im_k=images[1], mask_q=mask_q)
        # return loss and additional info (as list)
        return loss, [loss_moco, acc_moco, loss_plmoco, acc_plmoco]
# unpack data # compute output # return loss and additional info (as list)