def __init__(self,
trainLoader,
valLoader,
nbCls,
nClsEpisode,
nFeat,
outDir,
milestones=[50],
cuda=False):
self.trainLoader = trainLoader
self.valLoader = valLoader
self.outDir = outDir
self.milestones = milestones
if not os.path.isdir(self.outDir):
os.mkdir(self.outDir)
# Define model
self.netFeat, nFeat = get_featnet('WRN_28_10')
self.netClassifier = ClassifierTrain(nbCls)
self.netClassifierVal = ClassifierEval(nClsEpisode, nFeat)
# GPU setting
self.device = torch.device('cuda' if cuda else 'cpu')
if cuda:
self.netFeat.cuda()
self.netClassifier.cuda()
self.netClassifierVal.cuda()
self.criterion = nn.CrossEntropyLoss()
self.bestAcc = 0
valTransform,
args.cuda)
testLoader = EpisodeSampler(imgDir = testDir,
nClsEpisode = args.nClsEpisode,
nSupport = args.nSupport,
nQuery = args.nQuery,
transform = valTransform,
useGPU = args.cuda,
inputW = inputW,
inputH = inputH)
#############################################################################################
## Networks
netFeat, args.nFeat = get_featnet(args.architecture, inputW, inputH)
netFeat = netFeat.to(device)
netSIB = ClassifierSIB(args.nClsEpisode, args.nFeat, args.nStep)
netSIB = netSIB.to(device)
## Optimizer
optimizer = torch.optim.SGD(itertools.chain(*[netSIB.parameters(),]),
args.lr,
momentum=args.momentum,
weight_decay=args.weightDecay,
nesterov=True)
## Loss
criterion = nn.CrossEntropyLoss()
## Algorithm class