%% 画像データのロード
imdsTrain = imageDatastore('./trainingdata/','IncludeSubfolders',true, ...
    'LabelSource','foldernames');
imdsValidation = imageDatastore('./validationdata/','IncludeSubfolders',true, ...
    'LabelSource','foldernames');

%% SqueezeNet の読み込み
net = squeezenet;
inputSize = net.Layers(1).InputSize;

%% 最終レイヤの置換
lgraph = layerGraph(net); 
numClasses = numel(categories(imdsTrain.Labels));
newConvLayer =  convolution2dLayer([1, 1],numClasses, ...
    'WeightLearnRateFactor',10,'BiasLearnRateFactor',10,"Name",'new_conv');
lgraph = replaceLayer(lgraph,'conv10',newConvLayer);
newClassificatonLayer = classificationLayer('Name','new_classoutput');
lgraph = replaceLayer(lgraph,'ClassificationLayer_predictions',newClassificatonLayer);

%% Data augmentation
augimdsTrain = augmentedImageDatastore(inputSize(1:2),imdsTrain);
augimdsValidation = augmentedImageDatastore(inputSize(1:2),imdsValidation);

%% 学習パラメータ設定
options = trainingOptions('sgdm', ...
    'MiniBatchSize',11, ...
    'MaxEpochs',7, ...
    'InitialLearnRate',2e-4, ...
    'Shuffle','every-epoch', ...
    'ValidationData',augimdsValidation, ...
    'ValidationFrequency',3, ...
    'Verbose',false, ...
    'Plots','training-progress');

%% 学習の実行（別ウィンドウが開き，学習の進行状況が表示される）
netTransfer = trainNetwork(augimdsTrain,lgraph,options);
save('trainednetwork.mat','netTransfer') % 学習済みnetworkの保存

%% Copyright 2022 Mutsumi Saito