def loadPreTrained(model: BaseNet, path: str, logger: HtmlLogger) -> dict:
optimizerStateDict = None
if path is not None:
if exists(path):
# load checkpoint
checkpoint = loadModel(
path, map_location=lambda storage, loc: storage.cuda())
# load weights
model.loadPreTrained(checkpoint['state_dict'])
# # load optimizer state dict
# optimizerStateDict = checkpoint['optimizer']
# add info rows about checkpoint
loggerRows = []
loggerRows.append(['Path', '{}'.format(path)])
validationAccRows = [[
'Ratio', 'Accuracy'
]] + HtmlLogger.dictToRows(checkpoint['best_prec1'],
nElementPerRow=1)
loggerRows.append(['Validation accuracy', validationAccRows])
# set optimizer table row
optimizerRow = HtmlLogger.dictToRows(
optimizerStateDict, nElementPerRow=3
) if optimizerStateDict else optimizerStateDict
loggerRows.append(['Optimizer', optimizerRow])
logger.addInfoTable('Pre-trained model', loggerRows)
else:
raise ValueError(
'Failed to load pre-trained from [{}], path does not exists'
.format(path))
return optimizerStateDict
def train(self, trainFolderName='init_weights_train'):
args = self.getArgs()
# create train folder
folderPath = '{}/{}'.format(self.getTrainFolderPath(), trainFolderName)
if not exists(folderPath):
makedirs(folderPath)
# init optimizer
optimizer = self._initOptimizer()
# init scheduler
scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=0.95,
patience=args.weights_patience,
min_lr=args.learning_rate_min)
epoch = 0
trainLoggerFlag = True
while not self.stopCondition(epoch):
# update epoch number
epoch += 1
# init train logger
trainLogger = None
if trainLoggerFlag:
trainLogger = HtmlLogger(folderPath, epoch)
trainLogger.addInfoTable('Learning rates', [[
'optimizer_lr', self.formats[self.lrKey](
optimizer.param_groups[0]['lr'])
]])
# update train logger condition for next epoch
trainLoggerFlag = ((epoch + 1) % args.logInterval) == 0
# set loggers dictionary
loggersDict = {self.trainLoggerKey: trainLogger}
print('========== Epoch:[{}] =============='.format(epoch))
# train
trainData = self.weightsEpoch(optimizer, epoch, loggersDict)
# validation
validData = self.inferEpoch(epoch, loggersDict)
# update scheduler
scheduler.step(self.schedulerMetric(validData.lossDict()))
self.postEpoch(epoch, optimizer, trainData, validData)
self.postTrain()
def printToFile(self, saveFolder):
fileName = 'model'
logger = HtmlLogger(saveFolder, fileName)
if exists(logger.fullPath):
return
logger.setMaxTableCellLength(1000)
layerIdxKey = 'Layer#'
nFiltersKey = 'Filters#'
widthsKey = 'Width'
layerArchKey = 'Layer Architecture'
logger.createDataTable(
'Model architecture',
[layerIdxKey, nFiltersKey, widthsKey, layerArchKey])
for layerIdx, layer in enumerate(self._layers.flops()):
widths = layer.widthList()
dataRow = {
layerIdxKey: layerIdx,
nFiltersKey: layer.outputChannels(),
widthsKey: [widths],
layerArchKey: layer
}
logger.addDataRow(dataRow)
layerIdx += 1
# log additional layers, like Linear, MaxPool2d, AvgPool2d
for layer in self.additionalLayersToLog():
dataRow = {layerIdxKey: layerIdx, layerArchKey: layer}
logger.addDataRow(dataRow)
layerIdx += 1
# log layers alphas distribution
self.logTopAlphas(len(widths),
loggerFuncs=[
lambda k, rows: logger.addInfoTable(
self._alphasDistributionKey, rows)
],
logLayer=True)
# reset table max cell length
logger.resetMaxTableCellLength()