mirror of https://github.com/sthalles/SimCLR.git
16 lines
346 B
Python
16 lines
346 B
Python
from simclr import SimCLR
|
|
import yaml
|
|
from data_aug.dataset_wrapper import DataSetWrapper
|
|
|
|
|
|
def main():
|
|
config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader)
|
|
dataset = DataSetWrapper(config['batch_size'], **config['dataset'])
|
|
|
|
simclr = SimCLR(dataset, config)
|
|
simclr.train()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|